給一個樹,問有多少三元組滿足兩兩距離相等。n<=100000
長鏈剖分應用之一:o(n)統計以深度爲下標的可合併子樹信息
在當前節點,令f(i)表示相對深度爲i的節點個數,g(i)表示在子樹外離當前點距離爲i的點可以和子樹內多少對點組成答案。
每次新來一個兒子,枚舉長度,用當前g和兒子f以及當前f和兒子g更新一遍答案,然後用兩邊的f來更新g,再將兒子的f和g推入當前f和g。注意順序問題。
初始時直接繼承長兒子的信息,f數組相當於整體右移一位(指針左移一位),g數組相當於整體右移一位。一定要注意繼承時也要考慮這個點自己和長兒子所在子樹內組成的答案數(也就是繼承後的g[0])
在dfs時給每條長鏈底端分配相當於長鏈長度的內存,注意g指針是每次右移的,右移了len次後還需要有len的空間,因此應分配2*len。
就可以o(n)解決這道題了。
#include<cstdio>
#define gm 100005
using namespace std;
typedef long long ll;
inline ll* __alloc(size_t size)
{
static ll pool[gm<<2];
static ll* ptr=pool;
ll* res=ptr;ptr+=size;
return res;
}
struct e
{
int t;
e *n;
e(int t,e *n):t(t),n(n){}
}*p[gm];
int dep[gm],son[gm],len[gm],fa[gm];
ll *F[gm],*G[gm];
void dfs(int x)
{
son[x]=x;
for(e *i=p[x];i;i=i->n)
{
if(i->t==fa[x]) continue;
fa[i->t]=x;
dep[i->t]=dep[x]+1;
dfs(i->t);
if(dep[son[i->t]]>dep[son[x]]) son[x]=son[i->t];
}
len[x]=dep[son[x]]-dep[x]+1;
for(e *i=p[x];i;i=i->n)
{
if(i->t==fa[x]) continue;
if(son[i->t]!=son[x])
{
int y=son[i->t];
F[y]=__alloc(len[i->t])+len[i->t]-1;
G[y]=__alloc(len[i->t]<<1);
}
}
if(x==1)
{
int y=son[x];
F[y]=__alloc(len[x])+len[x]-1;
G[y]=__alloc(len[x]<<1);
}
}
ll ans=0;
void DP(int x)
{
ll *&f=F[x],*&g=G[x];
for(e *i=p[x];i;i=i->n)
{
if(i->t==fa[x]) continue;
DP(i->t);
if(son[x]==son[i->t])
{
f=F[i->t]-1;
g=G[i->t]+1;
}
}
ans+=g[0];
++f[0];
for(e *i=p[x];i;i=i->n)
{
if(i->t==fa[x]||son[i->t]==son[x]) continue;
ll *fs=F[i->t],*gs=G[i->t];
ans+=fs[0]*g[1];
for(int w=1;w<len[i->t];++w)
ans+=fs[w]*g[w+1]+gs[w]*f[w-1];
g[1]+=fs[0]*f[1];
f[1]+=fs[0];
for(int w=1;w<len[i->t];++w)
{
g[w+1]+=fs[w]*f[w+1];
g[w-1]+=gs[w];
f[w+1]+=fs[w];
}
}
}
int n;
int main()
{
scanf("%d",&n);
for(int i=1;i<n;++i)
{
int u,v;
scanf("%d%d",&u,&v);
p[u]=new e(v,p[u]);
p[v]=new e(u,p[v]);
}
dfs(1);
DP(1);
printf("%lld\n",ans);
return 0;
}