BZOJ3522&4543: [POI2014]Hotel加强版

题解见http://blog.csdn.net/neither_nor/article/details/51278993

不知道这是什么狗题，时间复杂度貌似是O(N)

#include<iostream>
#include<cstdio>
#include<cstring>
#define M 3000010
#define N 200010
using namespace std;
long long wa[M];
long long *f
,*g
,*p=wa;
long long to
,nxt
,pre
,cnt;
void ae(long long ff,long long tt)
{
cnt++;
to[cnt]=tt;
nxt[cnt]=pre[ff];
pre[ff]=cnt;
}
long long d
,fa
,mx
,yz
;
void build(long long x)
{
long long i,j;
yz[x]=x;
for(i=pre[x];i;i=nxt[i])
{
j=to[i];
if(j==fa[x]) continue;
fa[j]=x;
build(j);
if(d[j]>d[mx[x]])
{
yz[x]=yz[j];
mx[x]=j;
}
}
if(mx[x]==x) d[x]=0;
else d[x]=d[mx[x]]+1;
for(i=pre[x];i;i=nxt[i])
{
j=to[i];
if(j==fa[x]||(j==mx[x]&&x!=1)) continue;
p+=(d[j]+5);
f[yz[j]]=p;p+=5;
g[yz[j]]=p;
p+=2*d[j]+5;
}
}
long long ans;
void dfs(long long x)
{
long long i,j,k;
for(i=pre[x];i;i=nxt[i])
{
j=to[i];
if(j==fa[x]) continue;
dfs(j);
if(j==mx[x])
{
f[x]=f[j]-1;
g[x]=g[j]+1;
}
}
ans+=g[x][0];f[x][0]=1;
for(i=pre[x];i;i=nxt[i])
{
j=to[i];
if(j==fa[x]||j==mx[x]) continue;
ans+=g[x][1];
for(k=1;k<=d[j];k++)
{
ans+=g[j][k]*f[x][k-1];
ans+=f[j][k]*g[x][k+1];
}
for(k=0;k<=d[j];k++)
{
g[x][k-1]+=g[j][k];
g[x][k+1]+=f[j][k]*f[x][k+1];
f[x][k+1]+=f[j][k];
}
}
}
long long n;
int main()
{
scanf("%lld",&n);
long long i,j,x,y;
for(i=1;i<n;i++)
{
scanf("%lld%lld",&x,&y);
ae(x,y);ae(y,x);
}
build(1);
if(n==1) {puts("0");return 0;}
dfs(1);
printf("%lld\n",ans);
}