洛谷傳送門

為什麼總有人說這是長鏈剖分板子題？

Solution

題意非常簡潔，這讓我少了轉化題意這一步。

我們考慮什麼樣的三個點在樹上滿足兩兩之間距離 \(d\) 相等：

1. 對於某一個點，它的子樹內以它為LCA，距它 \(d\) 的三個點
2. 對於某一個點，它的 \(d\) 級祖先以及子樹內兩個以它為LCA，距它 \(d\) 的點

那麼我們選擇用樹形DP。

對於情況1，設 \(f_{i,j}\) 表示以 \(i\) 為根的子樹中，距 \(i\) 距離為 \(j\) 的點數。

那麼情況2怎麼表示比較好？看題解發現可以用上面提到的 某一個點 ，即中間點 。

設 \(g_{i,j}\) 表示以 \(i\)

那麼可以得到最樸素的轉移：

這個轉移有的地方比較好想，但是考慮周全還是要謹慎思考的。至於這裡面每個轉移的含義，畫個圖會恍然大悟的。

講完了

因為是最樸素的，在會列舉子樹的兩個點的轉移中，複雜度顯然是會爆炸的。

來到第一個優化——利用字首和的思想

考慮將 \(i\) 的一個新兒子 \(u\) 加入樹中對 \(ans\) 產生的貢獻，應該是 \(ans+=g_{i,j}\times f_{u,j-1}\) ，因為這個兒子只對之前加入的有貢獻（也可以自己模擬插入兒子算算）

同理可以得到： \(g_{i,j}+=f_{i,j}\times f_{u,j-1}\)

那麼這五種轉移就都是 \(O(n)\) 轉移的了，時間複雜度就變成了 \(O(n^2)\) (´▽`ʃ♡ƪ)

完結

什麼，你說我上面提到了長鏈剖分？

這裡是第二次優化——長鏈剖分

仔細看最後兩個轉移： \(f_{i,j}+=\sum_{u\in son_i}f_{u,j-1},g_{i,j}+=\sum_{u\in son_i}g_{u,j+1}\)

這樣可能不明顯，如果只有一個兒子： \(f_{i,j}=f_{son,j-1},g_{i,j}=g_{son,j+1}\)

也就是說，如果只有一個兒子，我們是可以做到直接賦值的

用指標寫就是 \(f_i=f_{son}-1,g_{i}=g_{son}+1\)

如果是在樹上，我們可以進行長鏈剖分，讓重兒子直接賦值，輕兒子和原來一樣轉移

因為輕兒子都是重鏈的頂部，所以轉移複雜度就是重鏈長度

那麼全域性時間複雜度就是 \(O(\sum len)=O(n)\)

正式完結(≧∇≦)ﾉ

Code

#include<cmath>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#define ll long long

using namespace std;
const int N=100010;
int n;

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

struct edge{
    int to,nxt;
}e[N<<1];
int head[N],cnt;
inline void add(int u,int v){
    e[++cnt].to=v;
    e[cnt].nxt=head[u];
    head[u]=cnt;
}

int h[N],son[N];
void dfs1(int u,int fa){
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v==fa) continue;
        dfs1(v,u);
        h[u]=max(h[u],h[v]);
        if(h[v]>h[son[u]]) son[u]=v;
    }
    h[u]=h[son[u]]+1;
}

ll *f[N],*g[N],tmp[N<<2],*id=tmp,ans;//這裡是指標
void dfs2(int u,int fa){
    if(son[u]) f[son[u]]=f[u]+1,g[son[u]]=g[u]-1,dfs2(son[u],u);
    f[u][0]=1;ans+=g[u][0];
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v==fa||v==son[u]) continue;
        f[v]=id;id+=h[v]<<1;
        g[v]=id;id+=h[v]<<1;
        dfs2(v,u);
        for(int j=0;j<h[v];j++){
            if(j) ans+=f[u][j-1]*g[v][j];
            ans+=g[u][j+1]*f[v][j];
        }
        for(int j=0;j<h[v];j++){
            g[u][j+1]+=f[u][j+1]*f[v][j];
            if(j) g[u][j-1]+=g[v][j];
            f[u][j+1]+=f[v][j];
        }
    }
}

int main(){
    n=read();
    for(int i=1,u,v;i<n;i++){
        u=read();v=read();
        add(u,v);add(v,u);
    }
    dfs1(1,0);
    f[1]=id;id+=h[1]<<1;
    g[1]=id;id+=h[1]<<1;
    dfs2(1,0);
    printf("%lld\n",ans);
    return 0;
}