点分治板题。
说一下几个要注意的:
选重心的时候一开始要把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';
}