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;
}