BZOJ4543 POI2014 Hotel加强版

---

Description

同OJ3522
数据范围：n<=100000

Sample Input

7
1 2
5 7
2 5
2 3
5 6
4 5

Sample Output

5

---

---

#include<bits/stdc++.h>

using namespace std;

#define LL long long

#define N 100010

LL pool[N<<4];

LL* top=pool;

LL* get(int len){LL* t=top;top+=len;return t;}

LL *f[N],*g[N];

int n;

LL ans=0;

vector<int> p[N];

int dep[N],lson[N];

void dfs1(int u,int fa){

    dep[u]=0;lson[u]=0;

    for(int i=0;i<p[u].size();i++){

        int v=p[u][i];

        if(v==fa)continue;

        dfs1(v,u);

        dep[u]=max(dep[u],dep[v]+1);

        if(dep[v]>dep[lson[u]])lson[u]=v;

    }

}

void dfs2(int u,int fa,int& maxlen,int blank){

    maxlen=max(maxlen,dep[u]);

    if(lson[u]){

        dfs2(lson[u],u,maxlen,blank+1);

        ans+=g[lson[u]][1];

        f[u]=f[lson[u]]-1;

        f[u][0]=1;

        g[u]=g[lson[u]]+1;

    }else{

        f[u]=get(maxlen+5+blank)+blank;

        g[u]=get(maxlen+5+blank);

        f[u][0]=1;

    }

    for(int i=0;i<p[u].size();i++){

        int v=p[u][i],mxlen=0;

        if(v==fa||v==lson[u])continue;

        dfs2(v,u,mxlen,0);

        for(int j=0;j<dep[v];j++)ans+=f[u][j]*g[v][j+1];

        for(int j=1;j<=dep[v]+1;j++)ans+=g[u][j]*f[v][j-1];

        for(int j=1;j<=dep[v]+1;j++)g[u][j]+=f[u][j]*f[v][j-1];

        for(int j=0;j<=dep[v];j++)f[u][j+1]+=f[v][j];

        for(int j=1;j<=dep[v];j++)g[u][j-1]+=g[v][j];

    }

}

int main(){

    scanf("%d",&n);

    for(int i=1;i<n;i++){

        int u,v;

        scanf("%d%d",&u,&v);

        p[u].push_back(v);

        p[v].push_back(u);

    }

    int mxlen=0;

    dfs1(1,0);

    dfs2(1,0,mxlen,0);

    printf("%lld",ans);

    return 0;

}