luoguP3806 【模板】点分治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

说明

对于30%的数据n<=100

对于60%的数据n<=1000,m<=50

对于100%的数据n<=10000,m<=100,c<=1000,K<=10000000

---

套用点分治的模板。

设num[k]为树中长度为k的路径的出现次数。

对于一个点no，可以dfs遍历一遍以它为根的子树中、以它开头向下的链长度。

对于每一对以no开头的链长度d[i],d[j]，num[d[i]+d[j]]++即可。

当然要记得容斥，因为要去掉d[i],d[j]表示的链来自no的同一个儿子。

设no与一个儿子相连边的长度为dis，在递归处理这个儿子的时候，对儿子的每一对d[i]，d[j]，做num[d[i]+d[j]+dis*2]--即可。

 #include<cstdio>
 #include<cstring>
 #include<iostream>
 #define rint register int
 using namespace std;

 int read(){
     char ch;
     int re=;
     bool flag=;
     while((ch=getchar())!='-'&&(ch<''||ch>''));
     ch=='-'?flag=:re=ch-'';
     while((ch=getchar())>=''&&ch<='')  re=re*+ch-'';
     return flag?-re:re;
 }

 struct Edge{
     int to,nxt,w;
     Edge(int to=,int nxt=,int w=):
         to(to),nxt(nxt),w(w){}
 };

 const int maxn=;

 int n,m,cnt=,sum,tot,root;
 int head[maxn],son[maxn],F[maxn],num[],d[maxn];
 bool vis[maxn];
 Edge E[maxn<<];

 inline void a_ed(int from,int to,int w){
     E[++cnt]=Edge(to,head[from],w);
     head[from]=cnt;
     E[++cnt]=Edge(from,head[to],w);
     head[to]=cnt;
 }

 void init(){
     n=read();  m=read();
     for(rint i=,from,to,w;i<n;i++){
         from=read();  to=read();  w=read();
         a_ed(from,to,w);
     }
 }

 void getroot(int no,int fa){
     son[no]=;  F[no]=;
     for(rint e=head[no];e;e=E[e].nxt){
         int nt=E[e].to;
         if(nt==fa||vis[nt])  continue;
         getroot(nt,no);
         son[no]+=son[nt];
         F[no]=max(F[no],son[nt]);
     }
     F[no]=max(F[no],sum-son[no]);
     if(F[no]<F[root])  root=no;
 }

 void getdeep(int no,int fa,int dd){
     d[tot++]=dd;
     for(rint e=head[no];e;e=E[e].nxt){
         int nt=E[e].to;
         if(nt==fa||vis[nt])  continue;
         getdeep(nt,no,dd+E[e].w);
     }
 }

 void calc(int no,bool opt,int p){
     tot=;
     getdeep(no,,);
     for(int i=;i<tot;i++)
         for(int j=;j<tot;j++)
             if(opt)  num[d[i]+d[j]]++;
             else  num[d[i]+d[j]+p]--;
 }

 void solve(int no){
     vis[no]=;
     calc(no,,);
     for(rint e=head[no];e;e=E[e].nxt){
         int nt=E[e].to;
         if(vis[nt])  continue;
         calc(nt,,E[e].w<<);
         sum=son[nt];  root=;
         getroot(nt,);
         solve(nt);
     }
 }

 int main(){
     init();
     sum=F[root=]=n;
     getroot(,);
     solve(root);
     for(rint i=,q;i<m;i++){
         q=read();
         if(num[q])  puts("AYE");
         else  puts("NAY");
     }
     return ;
 }