[bzoj4543] [POI2014]Hotel加强版

题目大意

给定一棵树，n个节点，问有多少个三元组(x,y,z)(x < y < z)，满足这三个点在树上距离两两相等。

n≤100000

分析

你可以想到一个n方的dp：设f[i][j]表示i为根的子树中，与i距离为j的节点有多少个。g[i][j]表示i为根的子树中，有多少个二元组(x,y)(x < y)满足：设d表示它们到lca的距离都为d，它们的lca到i的距离为d-j。

那么做到i节点时，先递归完每个儿子，然后合并每个子树。每次合并的时候，答案三元组中有1个、2个节点在当前被合并的子树中的答案加到总答案里、把g[x][1]加到总答案里，然后更新f和g。

如何跑过100000的数据呢？

考虑长链剖分。遍历完最深的子树后，用指针（我用数组模拟链表）操作O(1)复制该子树的f,g数组。然后再合并其它儿子，这时暴力合并。这样的复杂度是∑子树i的最大深度−i的深度。总复杂度是O(n)的

#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

const int N=1e5+5,M=N*2;

typedef long long LL;

int n,tot,h
,e[M],nxt[M],dep
,mx
,pf[M],pg[M],nf[M],ng[M],hf
,hg
,ef
,eg
,sf,sg;

LL f[M],g[M];

LL ans;

char c;

int read()
{
int x=0,sig=1;
for (c=getchar();c<'0' || c>'9';c=getchar()) if (c=='-') sig=-1;
for (;c>='0' && c<='9';c=getchar()) x=x*10+c-48;
return x*sig;
}

void Add(int x,int y)
{
e[++tot]=y; nxt[tot]=h[x]; h[x]=tot;
}

void Init(int x,int y)
{
mx[x]=dep[x]=dep[y]+1;
for (int i=h[x];i;i=nxt[i]) if (e[i]!=y)
{
Init(e[i],x); mx[x]=max(mx[x],mx[e[i]]);
}
}

void dfs(int x,int y)
{
int t=0,i;
for (i=h[x];i;i=nxt[i]) if (e[i]!=y && (!t || mx[t]<mx[e[i]])) t=e[i];
if (!t)
{
ef[x]=hf[x]=++sf; eg[x]=hg[x]=++sg; f[sf]=1; return;
}
dfs(t,x);
hf[x]=hf[t]; hg[x]=hg[t]; ef[x]=ef[t]; eg[x]=eg[t];
pf[hf[x]]=++sf; nf[sf]=hf[x]; f[sf]=1; hf[x]=sf;
ng[eg[x]]=++sg; pg[sg]=eg[x]; eg[x]=sg; ng[eg[x]]=++sg; pg[sg]=eg[x]; eg[x]=sg;
hg[x]=ng[hg[x]]; pg[hg[x]]=0;
ans+=g[hg[x]];
int xf,xg,yf,yg;
for (i=h[x];i;i=nxt[i]) if (e[i]!=y && e[i]!=t)
{
dfs(e[i],x);
for (xf=hf[x],xg=hg[x],yf=hf[e[i]],yg=hg[e[i]];yf>0 && yg>0;xf=nf[xf],yf=nf[yf],xg=ng[xg],yg=ng[yg])
{
if (pf[xf]>0) ans+=f[pf[xf]]*g[yg];
if (ng[xg]>0) ans+=g[ng[xg]]*f[yf];
}
for (xf=hf[x],xg=hg[x],yf=hf[e[i]],yg=hg[e[i]];yf>0 && yg>0;xf=nf[xf],yf=nf[yf],xg=ng[xg],yg=ng[yg])
{
if (pg[xg]>0) g[pg[xg]]+=g[yg];
if (ng[xg]>0) g[ng[xg]]+=f[nf[xf]]*f[yf];
if (nf[xf]>0) f[nf[xf]]+=f[yf];
}
}
}

int main()
{
n=read();
for (int i=1;i<n;i++)
{
int x=read(),y=read();
Add(x,y); Add(y,x);
}
Init(1,0);
dfs(1,0);
printf("%lld\n",ans);
return 0;
}