Codeforces 161D Distance in Tree（点分治/树形dp）

题意：给出一个树，求有多少对节点间的距离等于k，n<5e4, k<500

题解：解法一：点分治，k很小，在拿一个桶统计一下，注意统计的时候要判断时候和本身相同，相同要减一

#include <bits/stdc++.h>
#define IO_read ios::sync_with_stdio(false);cin.tie(0)
#define fre freopen("in.txt", "r", stdin)
#define _for(i,a,b) for(int i=a; i< b; i++)
#define _rep(i,a,b) for(int i=a; i<=b; i++)
#define inf 0x3f3f3f3f
#define lowbit(a) ((a)&-(a))
using namespace std;
typedef long long ll;
template <class T>
void read(T &x)
{
    char c; bool op=0;
    while(c=getchar(), c<'0'||c>'9') if(c=='-') op=1;
    x=c-'0';
    while(c=getchar(), c>='0'&&c<='9') x=x*10+c-'0';
    if(op) x=-x;
}
template <class T>
void write(T x)
{
    if(x<0) putchar('-'), x=-x;
    if(x>=10) write(x/10);
    putchar('0'+x%10);
}

const int maxn=5e4+5;
int head[maxn], tot;
struct Edge{
    int to, next, val;
}edge[maxn*2];

void addedge(int u, int v, int w)
{
    edge[++tot]={v, head[u], w};
    head[u]=tot;
}

int n, k;
ll ans;
int sz[maxn], vis[maxn], bucket[505], max_son, rt, Size;
ll dis[maxn], q[maxn];
int l, r;

void get_rt(int u, int fa)
{
    sz[u]=1;
    int num=0;
    for(int i=head[u]; i!=-1; i=edge[i].next)
    {
        int v=edge[i].to;
        if(vis[v] || v==fa) continue;
        get_rt(v, u);
        sz[u]+=sz[v];
        num=max(num, sz[v]);
    }
    num=max(num, Size-sz[u]);
    if(num<max_son) max_son=num, rt=u;
}

void get_dis(int u, int fa)
{
    q[++r]=dis[u];
    for(int i=head[u]; i!=-1; i=edge[i].next)
    {
        int v=edge[i].to, w=edge[i].val;
        if(vis[v] || v==fa) continue;
        dis[v]=dis[u]+w;
        get_dis(v, u);
    }
}

ll solve(int u, int val)
{
    l=1, r=0, dis[u]=val;
    get_dis(u, 0);
    ll res=0;
    sort(q+1, q+1+r);
    memset(bucket, 0, sizeof(bucket));
    for(int i=1; i<=r&&q[i]<=k; i++)
        bucket[q[i]]++;
    for(int i=1; i<=r&&q[i]<=k; i++)
    {
        if(k-q[i]==q[i]) res+=bucket[k-q[i]]-1;
        else res+=bucket[k-q[i]];
    }
    /*
    for(int i=1; i<=r; i++) printf("%d ", q[i]);
    printf("
");
    printf("res=%d, k=%d
", res);
    */
    res=res/2;
    return res;
}

void divide(int u)
{
    ans+=solve(u, 0);
    vis[u]=true;
    for(int i=head[u]; i!=-1; i=edge[i].next)
    {
        int v=edge[i].to, w=edge[i].val;
        if(vis[v]) continue;
        ans-=solve(v, w);  //容斥
        Size=sz[v], max_son=inf;
        get_rt(v, 0);
        divide(rt);
    }
}

int main()
{
    IO_read;
    //fre;
    cin>>n>>k;
    memset(head, -1, sizeof(head));
    for(int i=1,u,v; i<n; i++)
    {
        cin>>u>>v;
        addedge(u, v, 1), addedge(v, u, 1);
    }
    Size=n, max_son=inf;
    get_rt(1, 0);
    divide(rt);
    cout<<ans<<"
";
}

　　解法二：树形dp，待解决