Codeforces 161D Distance in Tree(点分治)

题意：求distance(u,v)==k的点对数。

思路：统计距离root距离之和为k的对数-距离root儿子节点中距离儿子节点距离之和为k-2的对数。
因为每对算了两边除二即可。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn = 1e6+5;
const int INF  = 0x3f3f3f3f;
int tot,head[maxn];
struct node{
    int v,w,nxt;
}e[maxn];
void add(int u,int v,int w){
    e[tot].v = v;
    e[tot].w = w;
    e[tot].nxt = head[u];
    head[u]=tot++;
}
void init(){
    tot = 0;
    memset(head,-1,sizeof(head));
}
int n,k;
ll ans;
int rt,all,num;
int sz[maxn],son[maxn],vis[maxn],dis[maxn];
 
void getroot(int u,int fa){
    sz[u]  = 1;
    son[u] = 0;
    for(int i=head[u];i!=-1;i=e[i].nxt){
        int v = e[i].v;
        if(v==fa || vis[v]) continue;
        getroot(v,u);
        sz[u]  += sz[v];
        son[u] = max(son[u],sz[v]);
    }
    son[u] = max(son[u],all-sz[u]);
    if(son[rt] > son[u] || son[rt]==0) rt = u;
}
void getdis(int u,int fa,int len,int dist){
    dis[++num] = dist;
    if(dist >= len) return ;
    for(int i=head[u];i!=-1;i=e[i].nxt){
        int v = e[i].v;
        if(v==fa || vis[v]) continue;
        getdis(v,u,len,dist+e[i].w);
    }
}
void cal(int u,int fa,int len,int op){
    if(len < 0) return ;
    num = 0;
    getdis(u,fa,len,0);
    ll res = 0;
    int tmp[505]{};
    for(int i=1;i<=num;i++) tmp[dis[i]]++;
    for(int i=1;i<=num;i++)  res += tmp[len-dis[i]];
    ans += 1ll*op*res;
}
void solve(int u){
    rt = 0;getroot(u,0);u = rt;
    cal(u,0,k,1);
    vis[u] = 1;
    for(int i=head[u];i!=-1;i=e[i].nxt){
        int v = e[i].v;
        if(vis[v]) continue;
        cal(v,u,k-2,-1);
        all = sz[v];
        solve(v);
    }
}
int main(){
    //freopen("in.txt","r",stdin);
    init();
    scanf("%d%d",&n,&k);
    for(int i=1;i<n;i++){
        int u,v;
        scanf("%d%d",&u,&v);
        add(u,v,1);add(v,u,1);
    }
    all = n;
    solve(1);
    printf("%lld\n",ans/2);
    return 0;
}