[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数组。然后再合并其它儿子，这时暴力合并。这样的复杂度是$\sum 子树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[N],e[M],nxt[M],dep[N],mx[N],pf[M],pg[M],nf[M],ng[M],hf[N],hg[N],ef[N],eg[N],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;
}