BZOJ 4543 POI2014 hotel

Problem

Solution

Code

#include <cstdio>
using namespace std;
typedef long long ll;
const int maxn=100010;
struct data{int v,nxt;}edge[maxn<<1];
int n,p,head[maxn],deep[maxn],mx[maxn];
ll ans,tmp[maxn*5];
ll *ptr=tmp+1,*f[maxn],*g[maxn];
inline void insert(int u,int v)
{
edge[++p]=(data){v,head[u]};head[u]=p;
edge[++p]=(data){u,head[v]};head[v]=p;
}
void dfs(int x,int pre)
{
mx[x]=x;
for(int i=head[x];i;i=edge[i].nxt)
if(edge[i].v!=pre)
{
deep[edge[i].v]=deep[x]+1;
dfs(edge[i].v,x);
if(deep[mx[x]]<deep[mx[edge[i].v]]) mx[x]=mx[edge[i].v];
}
for(int i=head[x];i;i=edge[i].nxt)
if(edge[i].v!=pre&&(mx[edge[i].v]!=mx[x]||x==1))
{

4000
int v=mx[edge[i].v];
ptr+=deep[v]-deep[x]+1;
f[v]=ptr;g[v]=(ptr+=1);
ptr+=(deep[v]-deep[x])<<1|1;
}
}
void dp(int x,int pre)
{
for(int i=head[x];i;i=edge[i].nxt)
if(edge[i].v!=pre)
{
dp(edge[i].v,x);
if(mx[x]==mx[edge[i].v]) f[x]=f[edge[i].v]-1,g[x]=g[edge[i].v]+1;
}
ans+=g[x][0];f[x][0]=1;
for(int i=head[x];i;i=edge[i].nxt)
if(edge[i].v!=pre&&mx[x]!=mx[edge[i].v])
{
int lim=deep[mx[edge[i].v]]-deep[x];
for(int j=0;j<=lim;j++)
ans+=f[x][j-1]*g[edge[i].v][j]+g[x][j+1]*f[edge[i].v][j];
for(int j=0;j<=lim;j++)
{
g[x][j-1]+=g[edge[i].v][j];
g[x][j+1]+=f[x][j+1]*f[edge[i].v][j];
f[x][j+1]+=f[edge[i].v][j];
}
}
}
void input()
{
scanf("%d",&n);
for(int i=1,x,y;i<n;i++)
{
scanf("%d%d",&x,&y);
insert(x,y);
}
}
int main()
{
#ifndef ONLINE_JUDGE
freopen("in.txt","r",stdin);
#endif
input();
deep[1]=1;
dfs(1,0);
dp(1,0);
printf("%lld\n",ans);
return 0;
}