Distance in Tree（点分治）

题目

 https://codeforces.com/problemset/problem/161/D

题意

给你n个点的树 问相距为K的有多少对

思路

点分治模板

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
const int inf = 10000000;
const int maxn = 50003;
int n,m,k;
struct node
{
    int v,nxt;
}E[maxn<<1];
int tot,head[maxn];
int maxp[maxn],siz[maxn];//当前节点最大子树节点数
int dis[maxn],rem[maxn],judge[maxn]; //根节点当当前节点长度 当前根节点到某子树所以子节点长度 当前根节点之前所以子树出现的长度
int vis[maxn];//当前节点是否已被分治
int test[maxn],q[maxn];//记录答案 记录出现的长度
int sum,rt;
ll ans;
void add(int u,int v )
{
    E[++tot].nxt = head[u];
    E[tot].v = v;
    head[u] = tot;
}

void getrt(int u,int pa) //查询重心
{
    siz[u] = 1;
    maxp[u] = 0;
    for(int i = head[u];i != -1;i = E[i].nxt)
    {
        int v = E[i].v;
        if(v == pa||vis[v]) continue;
        getrt(v,u);
        siz[u] += siz[v];
        maxp[u] = max(maxp[u],siz[v]);
    }
    maxp[u] = max(maxp[u],sum-siz[u]);
    if(maxp[u] < maxp[rt]) rt = u;
}
void getdis(int u,int fa) //查询节点到当前根的长度
{
    if(dis[u] <= k)
    {
        rem[dis[u]]++;
    }
    for(int i = head[u];i != -1;i=E[i].nxt)
    {
        int v = E[i].v;
        if(v == fa||vis[v]) continue;
        dis[v] = dis[u]+1;
        getdis(v,u);
    }
}

void calc(int u) //当前根节点出现的结果
{
    int p = 0;
    for(int i = head[u];i != -1;i=E[i].nxt)
    {
        int v = E[i].v;
        if(vis[v]) continue;
        memset(rem,0,sizeof(rem));
        dis[v] = 1;
        getdis(v,u);

        for(int j = 0;j <= k;j++)
        {
            ans += rem[j]*judge[k-j];
        }
        for(int j = 0;j <= k;j++)
        {
            judge[j] += rem[j];
        }
    }
    for(int i = 0;i <= k;i++)
    {
        judge[i] = 0;
    }
}

void solve(int u)
{
    vis[u] = judge[0] = 1;
    calc(u);
    for(int i = head[u];i != -1;i = E[i].nxt)
    {
        int v = E[i].v;
        if(vis[v]) continue;
        sum = siz[v],maxp[rt = 0] = inf;
        getrt(v,0),solve(rt);
    }
}
int  main()
{
    memset(head,-1,sizeof(head));
    tot = 0;
    scanf("%d%d",&n,&k);
    for(int i = 1;i < n;i++) // 建图
    {
        int u,v;
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    maxp[rt] = sum = n;
    getrt(1,0); //求重心

    memset(vis,0,sizeof(vis));
    ans = 0;
    solve(rt);//开始分治
    cout<<ans<<endl;
    return 0;
}