【点分治】Distance in tree

传送门

 

点分治板题。

说一下几个要注意的：

 

选重心的时候一开始要把mx[0]赋为n（总点数）。

处理一个子树时，把SIZ重置为当前子树的大小，root（重心）重置为0，tot重置为1（记录每个点到重心的距离）。

 

分治到一个点时，标记一下vis数组，表明这个点被分治了，后面的子问题就在只它的子树内解决。所以后面重新求重心，求距离的操作都要避开已经vis的点。

 

 

求两点之间距离小于等于k，考虑两点在当前重心不同子树内。对于在同一子树内的，递归处理。

把每个点到重心的距离处理出来之后排个序，然后枚举一个点x，如果它到重心距离对应为dis[l]（分治过后点的标号不一定是连续的。重新用一个参数tot记录，保证记录下来是连续的。具体看代码），那么我们就要找到距离为（k-dis[l]）的点。

考虑枚举左端点l。

用lower_bound找到大于等于（k-dis[l]）的第一个点，用upper_bound找到大于（k-dis[l]）的第一个点。

如果没有等于（k-dis[l]）的，lower_bound自然等于upper_bound。

所以（k-dis[l]）的个数就是upper_bound-lower_bound。

注意！这里二分查找的范围是[l,tot]！！因为前面如果有dis[k]=dis[l]，且k∈[1,l)，是不能算进去的！！

l的上界是tot-1，因为不能出现dis[tot]+dis[tot]=k，这样的话就相当于两个点是同一个点。

 

但是像左边这种情况是不行的，要去掉选重边的情况。就是算子树深度的时候强制把重边的长度算进去，然后把它的贡献减去就行了。具体见代码。

 

最后注意一下加边用的cnt不要和tot重了。

 

代码：

#include<bits/stdc++.h>
using namespace std;
const int maxn=5e4+10;
int Head[maxn],Next[maxn<<1],V[maxn<<1],W[maxn<<1],cnt=0;
int dep[maxn],siz[maxn],vis[maxn],mx[maxn],dis[maxn];
int n,a,b,c,k,SIZ,root=0,tot=0,ans=0;
inline void add(int u,int v,int w){++cnt,Next[cnt]=Head[u],V[cnt]=v,W[cnt]=w,Head[u]=cnt;}
void get_root(int u,int f){
	mx[u]=0,siz[u]=1;
	for(int i=Head[u];i;i=Next[i]) if((V[i]!=f)&&(!vis[V[i]]))
		get_root(V[i],u),mx[u]=max(mx[u],siz[V[i]]),siz[u]+=siz[V[i]];
	mx[u]=max(mx[u],SIZ-siz[u]);if(mx[u]<mx[root]) root=u;
}
void get_dis(int u,int f,int D){
	for(int i=Head[u];i;i=Next[i])if((V[i]!=f)&&(!vis[V[i]]))
		dis[++tot]=D+W[i],get_dis(V[i],u,dis[tot]);
}
int get_ans(int x,int D,int l=1,int ret=0){
	dis[tot=1]=D,get_dis(x,0,D),sort(dis+1,dis+tot+1);
	while(l<tot&&dis[l]+dis[tot]<k) ++l;
	while(l<tot&&(dis[l]<=k-dis[l]))
	ret+=((upper_bound(dis+l+1,dis+tot+1,k-dis[l])-dis)-(lower_bound(dis+l+1,dis+tot+1,k-dis[l])-dis)),++l;
	return ret;
}
void dfs(int u){
	vis[u]=1,ans+=get_ans(u,0);
	for(int i=Head[u];i;i=Next[i]) if(!vis[V[i]])
	ans-=get_ans(V[i],W[i]),SIZ=siz[V[i]],root=0,get_root(V[i],u),dfs(root);
}
int main(){
	scanf("%d%d",&n,&k),mx[0]=SIZ=n;
	for(int i=1;i<n;++i) scanf("%d%d",&a,&b),c=1,add(a,b,c),add(b,a,c);
	get_root(1,0),dfs(root),cout<<ans<<'\n';
}