P3806 【模板】点分治1

　　

题目背景

感谢hzwer的点分治互测。

题目描述

给定一棵有n个点的树

询问树上距离为k的点对是否存在。

输入格式

n,m 接下来n-1条边a,b,c描述a到b有一条长度为c的路径

接下来m行每行询问一个K

输出格式

对于每个K每行输出一个答案，存在输出“AYE”,否则输出”NAY”(不包含引号)

输入输出样例

输入 #1复制

2 1
1 2 2
2

输出 #1复制

AYE

除了work函数需要重写   其他的都比较模板

#include<bits/stdc++.h>
using namespace std;
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define repp(i,a,b) for(int i=(a);i>=(b);--i)
#define ll long long
#define see(x) (cerr<<(#x)<<'='<<(x)<<endl)
#define inf 0x3f3f3f3f
#define CLR(A,v)  memset(A,v,sizeof A)
//////////////////////////////////
const int N=1e5+10;
int n,m,ans[N],a[N],num,root,SIZE,siz[N],pos,head[N],maxx,f[N],vis[N];
struct Edge{int to,nex,w;}edge[N];
void add(int a,int b,int c)
{
    edge[++pos]=(Edge){b,head[a],c};
    head[a]=pos;
}
void getroot(int x,int fa)
{
    f[x]=0;siz[x]=1;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;if(vis[v]||v==fa)continue;
        getroot(v,x);
        f[x]=max(f[x],siz[v]);
        siz[x]+=siz[v];
    }
    f[x]=max(f[x],SIZE-siz[x]);
    if(f[x]<f[root])root=x;
}
struct node
{
    int id,v;
}dis[N];
void dfs(int x,int fa,int id,int d)
{
    dis[++num]=(node){id,d};
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v]||v==fa)continue;
        dfs(v,x,id,d+edge[i].w);
    }
}
int find1(int x)
{
    int ans=0,L=1,R=num;
    while(L<=R)
    {
        int mid=(L+R)>>1;
        if(dis[mid].v>=x)ans=mid,R=mid-1;
        else L=mid+1;
    }
    return ans;
}
void work(int x)
{
    num=0;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v])continue;
        dfs(v,x,v,edge[i].w);
    }
    dis[++num]=(node){0,0};
    sort(dis+1,dis+1+num,[](node a,node b){return a.v<b.v;});
    rep(i,1,m)
    {
        if(ans[i])continue;
        int L=1;
        while(L<num&&dis[L].v+dis[num].v<a[i])L++;
        while(L<num&&!ans[i])
        {
            if(2*dis[L].v>a[i])break;
            int pos=find1(a[i]-dis[L].v);
            while(L<=num&&dis[pos].id==dis[L].id&&dis[pos].v+dis[L].v==a[i])pos++;
            if(dis[pos].v+dis[L].v==a[i])ans[i]=1;
            L++;
        }
    }
}
void solve(int x)
{
    vis[x]=1;
    work(x);
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v])continue;
        root=0;
        siz[root]=SIZE=siz[v];
        getroot(v,0);
        solve(v);
    }
}
int main()
{
    cin>>n>>m;
    rep(i,1,n-1)
    {
        int a,b,c;cin>>a>>b>>c;
        add(a,b,c);add(b,a,c);
    }
    rep(i,1,m)cin>>a[i];

    f[root]=SIZE=n;
    getroot(1,0);
    solve(root);
    rep(i,1,m)
    if(ans[i])printf("AYE
");
    else printf("NAY
");
    return 0;
}

容斥写法

#include<bits/stdc++.h>
using namespace std;
#define ll long long 
const int N=2e5+6;
ll mod=1e9+7;
int n,m,cnt,dis[N],head[N],pos,siz[N],siz_tree,x,y,z,vis[N],root,maxson[N],ans[10000001];
struct Edge{int to,nex,v;}edge[N<<1];
void add(int a,int b,int c){edge[++pos]=(Edge){b,head[a],c};head[a]=pos;}

void getroot(int x,int fa)
{
    siz[x]=1;maxson[x]=0;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v]||v==fa)continue;
        getroot(v,x);siz[x]+=siz[v];
        maxson[x]=max(maxson[x],siz[v]);
    }
    maxson[x]=max(maxson[x],siz_tree-siz[x]);
    if(maxson[x]<maxson[root])root=x;
}

void getdis(int x,int fa,int nowdis)
{
    dis[++dis[0]]=nowdis;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(v==fa||vis[v])continue;
        getdis(v,x,nowdis+edge[i].v);
    }
}

void calc(int x,int d,int flag)
{
    dis[0]=0;
    getdis(x,0,d);
    for(int i=1;i<=dis[0];i++)
    for(int j=1;j<=dis[0];j++)if(i!=j&&dis[i]+dis[j]<=10000000)
    ans[dis[i]+dis[j]]+=flag;
}
void solve(int x)
{
    calc(x,0,1);
    vis[x]=1;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v])continue;
        calc(v,edge[i].v,-1);
        root=0;
        siz_tree=siz[v];
        getroot(v,x);
        solve(root);
    }
}
int main()
{
    cin>>n>>m;
    for(int i=1;i<n;i++)scanf("%d%d%d",&x,&y,&z),add(x,y,z),add(y,x,z);
    siz_tree=maxson[0]=n;
    root=0;
    getroot(1,0);
    solve(root);
    for(int i=1;i<=m;i++)scanf("%d",&x),printf("%s
",ans[x]>0?"AYE":"NAY");
}