[POI2014][BZOJ3522/4543]Hotel/[JZOJ5357]校門外的樹

題目大意

給定一棵n 個點的樹，求樹上兩兩距離相等的點三元組個數。

1≤n≤105

---

題目分析

考慮dp 。
令fx,i 表示x 爲根的子樹內，距離x 爲i 的點的個數；gx,i 表示以x 爲根的子樹中，到x 距離相等而且到lca 的距離比lca 到x 距離要大i 的點對個數（說白了就是那些可能的在x 子樹外的第三個點伸出了x 子樹i 的距離）。
然後在dp 各個子樹之前，我們有fx,0=1,ans+=gx,0 。
對於一個子樹y ，我們有

ansgx,i+1gx,i−1fx,i+1+=fx,igy,i+1+gx,i+1fy,i+=fx,i+1fy,i+=gy,i+=fy,i

這樣我們就可以做到O(n2) 的時間和空間複雜度。
怎麼將時空複雜度優化呢？注意到對於節點x ，設y 是我們第一個枚舉的子樹，那麼gx,i−1=gy,i ，fx,i+1=fy,i ，是數組位移一位的關係。
我們不妨對這棵樹長鏈剖分，然後將重兒子作爲第一次枚舉的子樹，使用類似指針的思路來做到O(1) 的數組位移，然後其餘的子樹直接暴力轉移。
這樣做時間複雜度是O(n) 的，因爲一個點x 所在的子樹被暴力轉移當且僅當x 是一條長鏈的頂端，而且轉移複雜度是O(depx) 的，也就是和長鏈長成正比的。因此總的轉移複雜度就是所有長鏈長的和，也就是O(n) 的。
至於空間複雜度，我們給每條長鏈頂端分配正比於長鏈長度的空間就好了，最後也是O(n) 的。
一些實現細節請讀者自行思考。

---

代碼實現

#include <iostream>
#include <cstdio>
#include <cctype>

using namespace std;

typedef long long LL;

int read()
{
    int x=0,f=1;
    char ch=getchar();
    while (!isdigit(ch)) f=ch=='-'?-1:f,ch=getchar();
    while (isdigit(ch)) x=x*10+ch-'0',ch=getchar();
    return x*f;
}

const int N=100005;
const int M=N<<1;
const int E=N<<1;

int last[N],hea[N],depth[N],lng[N],mxl[N],fa[N],fst[N],gst[N];
int tov[E],nxt[E];
LL f[N],g[M];
int n,tot,fcnt,gcnt;
LL ans;

void insert(int x,int y){tov[++tot]=y,nxt[tot]=last[x],last[x]=tot;}

void dfs(int x)
{
    mxl[x]=lng[x]=0;
    for (int i=last[x],y;i;i=nxt[i])
        if ((y=tov[i])!=fa[x])
        {
            depth[y]=depth[fa[y]=x]+1,dfs(y);
            if (mxl[x]<mxl[y]+1) mxl[x]=mxl[lng[x]=y]+1;
        }
}

void dp(int x,int top)
{
    if (x==top) fst[x]=fcnt,fcnt+=mxl[x]+1,gst[x]=gcnt,gcnt+=mxl[x]<<1|1;
    if (lng[x]) dp(lng[x],top);
    int fptr=fst[top]+depth[x]-depth[top],gptr=gst[top]+mxl[top]-depth[x]+depth[top];
    ++f[fptr],ans+=g[gptr];
    for (int i=last[x],y;i;i=nxt[i])
        if ((y=tov[i])!=fa[x]&&y!=lng[x])
        {
            dp(y,y);
            for (int j=0;j<mxl[y];++j) ans+=f[fptr+j]*g[gst[y]+mxl[y]+j+1];
            for (int j=0;j<=mxl[y]&&gptr+j+1<gst[top]+(mxl[top]<<1|1);++j) ans+=g[gptr+j+1]*f[fst[y]+j];
            for (int j=0;j<=mxl[y]&&gptr+j+1<gst[top]+(mxl[top]<<1|1)&&fptr+j+1<fst[top]+mxl[top]+1;++j) g[gptr+j+1]+=f[fptr+j+1]*f[fst[y]+j];
            for (int j=0;j<mxl[y];++j) g[gptr+j]+=g[gst[y]+mxl[y]+j+1];
            for (int j=0;j<=mxl[y]&&fptr+j+1<fst[top]+mxl[top]+1;++j) f[fptr+j+1]+=f[fst[y]+j];
        }
}

int main()
{
    freopen("tree.in","r",stdin),freopen("tree.out","w",stdout);
    n=read();
    for (int i=1,x,y;i<n;++i) x=read(),y=read(),insert(x,y),insert(y,x);
    depth[1]=1,dfs(1),fcnt=gcnt=1,dp(1,1),printf("%lld\n",ans);
    fclose(stdin),fclose(stdout);
    return 0;
}