codeforces 161D D. Distance in Tree (树分治)

题意：

问一颗树有多少个路径距离为k

思路：

是poj1741的简单版，在树分治的时候，对于某个重心只要维护某个长度的有多少即可，因为k很小，所以可以直接开个数组存

错误及反思：

中间加的过程忘了会爆int了。。。

代码：

#include<bits/stdc++.h>
#define fi first
#define se second
using namespace std;
const int N = 50100;
struct EDGE{
    int to,next;
}e[N*2];
int first[N],n,tot=0,si[N],maxn[N],num[501],k;
bool did[N];
inline int read()
{
    int x=0;char ch=getchar();
    while(ch<'0'||ch>'9')ch=getchar();
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x;
}
void addedge(int x,int y){
    e[tot].to=y;
    e[tot].next=first[x];
    first[x]=tot++;
    e[tot].to=x;
    e[tot].next=first[y];
    first[y]=tot++;
}
void dfs1(int now,int fa){
    si[now]=1;
    maxn[now]=0;
    for(int i=first[now];i!=-1;i=e[i].next)
        if(e[i].to!=fa&&!did[e[i].to]){
            dfs1(e[i].to,now);
            si[now]+=si[e[i].to];
            maxn[now]=max(maxn[now],si[e[i].to]);
        }
}
void dfs3(int now,int fa,int& root,int& num,int t){
    int MA=max(maxn[now],si[t]-si[now]);
    if(MA<num){
        num=MA;
        root=now;
    }
    for(int i=first[now];i!=-1;i=e[i].next)
        if(e[i].to!=fa&&!did[e[i].to])
            dfs3(e[i].to,now,root,num,t);
}
void dfs2(int now,int fa,int tlen){
    if(tlen>k) return ;
    for(int i=first[now];i!=-1;i=e[i].next)
        if(e[i].to!=fa&&!did[e[i].to])
            dfs2(e[i].to,now,tlen+1);
    num[tlen]++;
}
long long cal(int now,int d){
    memset(num,0,sizeof(num));
    dfs2(now,-1,d);
    long long ans=0;
    for(int i=0;i<=500;i++){
        if(i==k-i)
            ans+=1ll*num[i]*(num[i]-1);
        else
            ans+=1ll*num[i]*num[k-i];
    }
    return ans/2;
}
long long solve(int now){
    int root,num=1e9;
    dfs1(now,-1);
    dfs3(now,-1,root,num,now);
    long long ans=cal(root,0);
    did[root]=true;
    for(int i=first[root];i!=-1;i=e[i].next)
        if(!did[e[i].to]){
            ans-=cal(e[i].to,1);
            ans+=solve(e[i].to);
        }
    return ans;
}
int main(){
    memset(first,-1,sizeof(first));
    n=read();
    k=read();
    for(int i=0,u,v;i<n-1;i++){
        u=read(),v=read();
        addedge(u,v);
    }
    printf("%lld\n",solve(1));
}