CF161D Distance in Tree |点分治

Posted 2020-11-24 naruto-mzx

题目大意

输入点数为N一棵树

求树上长度恰好为K的路径个数

输入格式

第一行两个数字N,K,如题意

接下来的N?1行中,每行两个整数u,v表示一条树边(u,v)

输出格式

一个整数ans,如题意

---

直接套模板，没啥好说的

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<iostream>
#include<map>
using namespace std;
inline int read(){
    int x=0,f=1; char c=getchar();
    while(c<'0'||c>'9'){ if(c=='-')f=-1; c=getchar(); }
    while('0'<=c&&c<='9'){ x=x*10+c-'0'; c=getchar(); }
    return x*f;
}
const int N=5e4+10;
int nxt[N<<1],head[N],go[N<<1],tot;
inline void add(int u,int v){
    nxt[++tot]=head[u]; head[u]=tot; go[tot]=v;
    nxt[++tot]=head[v]; head[v]=tot; go[tot]=u; 
}
bool vis[N];
int n,k,size[N],sum,rt,maxp;
inline void getrt(int u,int fa){
    size[u]=1; int Max=0;
    for(int i=head[u];i;i=nxt[i]){
        int v=go[i];
        if(v==fa||vis[v])continue;
        size[u]+=size[v];
        Max=max(Max,size[v]);
    }
    Max=max(Max,sum-size[u]);
    if(Max<maxp)maxp=Max,rt=u;
}
int judge[N],rem[N],dis[N];
inline void getdis(int u,int fa){
    if(dis[u]>k)return;
    rem[++rem[0]]=dis[u];
    for(int i=head[u];i;i=nxt[i]){
        int v=go[i];
        if(v==fa||vis[v])continue;
        dis[v]=dis[u]+1;
        getdis(v,u);
    }
}
long long ans=0;
inline void calc(int u){
    for(int i=head[u];i;i=nxt[i]){
        int v=go[i];
        if(vis[v])continue;
        rem[0]=0; dis[v]=1; getdis(v,u);
        for(int i=1;i<=rem[0];i++)ans+=judge[k-rem[i]];
        for(int i=1;i<=rem[0];i++)judge[rem[i]]++;
    }
    for(int i=0;i<=n;i++){
        if(!judge[i])break;
        judge[i]=0;
    }
}
inline void solve(int u){
    judge[0]=vis[u]=1; calc(u);
    for(int i=head[u];i;i=nxt[i]){
        int v=go[i];
        if(vis[v])continue;
        sum=maxp=size[v]; getrt(v,0);
        solve(rt);
    }
}
signed main(){
    n=read(),k=read();
    for(int i=1;i<n;i++)add(read(),read());
    sum=maxp=n; getrt(1,0);
    solve(rt);
    cout<<(ans>>1)<<endl;
}