题解 CF161D 【Distance in Tree】

da32s1da · 2018-12-10 11:08:22 · 题解

在统计答案的时候，我和各位大佬的方法可能略有不同。

我们先统计小于等于k的数量，再减去小于k的数量。

于是类似P4178即可。

#include<cstdio>
#include<algorithm>
using namespace std;
const int N=5e4+50;
typedef long long LL;
int to[N<<1],las[N<<1],fir[N],ds[N<<1],cnt;
inline void add_edge(int u,int v,int a){
    to[++cnt]=v;las[cnt]=fir[u];fir[u]=cnt;ds[cnt]=a;
    to[++cnt]=u;las[cnt]=fir[v];fir[v]=cnt;ds[cnt]=a;
}
inline int max(int u,int v){return u>v?u:v;}
inline int min(int u,int v){return u<v?u:v;}
int n,k,x,y;
int f[N],dep[N],siz[N],vis[N],sth[N];
int Cnt,rot,sum;
LL ans;
void grot(int u,int v){
    siz[u]=1;f[u]=0;
    for(int i=fir[u];i;i=las[i])
    if(to[i]!=v&&!vis[to[i]]){
        grot(to[i],u);
        siz[u]+=siz[to[i]];
        f[u]=max(f[u],siz[to[i]]);
    }
    f[u]=max(f[u],sum-siz[u]);
    if(f[u]<f[rot])rot=u;
}
void gsth(int u,int v){
    sth[++Cnt]=dep[u];
    for(int i=fir[u];i;i=las[i])
    if(to[i]!=v&&!vis[to[i]]){
        dep[to[i]]=dep[u]+ds[i];
        gsth(to[i],u);
    }
}
int calc(int u,int dis){
    dep[u]=dis;Cnt=0;
    gsth(u,0);
    sort(sth+1,sth+Cnt+1);

    int l=1,r=Cnt,res=0;
    while(l<r)
    if(sth[l]+sth[r]<=k)res+=r-l,l++;
    else r--;
    //加上 长度 小于等于k 的

    l=1,r=Cnt;
    while(l<r)
    if(sth[l]+sth[r]<k)res-=r-l,l++;
    else r--;
    //减去 长度 小于k 的

    return res;
}
void solve(int u){
    ans+=calc(u,0);vis[u]=1;
    for(int i=fir[u];i;i=las[i])
    if(!vis[to[i]]){
        ans-=calc(to[i],ds[i]);
        rot=0;sum=siz[to[i]];
        grot(to[i],0);
        solve(rot);
    }
}
int main(){
    scanf("%d%d",&n,&k);
    for(int i=1;i<n;i++){
        scanf("%d%d",&x,&y);
        add_edge(x,y,1);
    }
    f[0]=sum=n;
    grot(1,0);
    solve(rot);
    printf("%I64d\n",ans);
}