题解：AT_abc359_g [ABC359G] Sum of Tree Distance

· · 题解

这题让我们求相同颜色（即 A_i）点之间互相的最短距离之和，由于颜色数量较多，很容易想到树上启发式合并。

设 \operatorname{LCA}(i,j) 表示点 i 与 j 在以 1 为根的有根树下的最近公共祖先，deep_i 表示 i 的深度，其中 deep_1=1。那么 i 与 j 的距离即为 deep_i+deep_j-2deep_{\operatorname{LCA}(i,j)}。

所有 A_i=a 对答案的贡献即为 \frac{1}{2}\sum\limits_{A_i=a}\sum\limits_{A_j=a}(deep_i+deep_j-2deep_{\operatorname{LCA}(i,j)})。

设以 t 为根的子树内满足 A_i=a 的点数量有 s_{t,a} 个。

不难发现前面的 \sum\limits_{A_i=a}\sum\limits_{A_j=a}deep_i+deep_j 非常好求，这个值就是：

2s_{1,a}\sum\limits_{A_i=a}deep_i

考虑用树上启发式合并求后面的值。不难发现我们可以把 deep_{\operatorname{LCA}(i,j)} 的图论意义解释为 \operatorname{LCA}(i,j) 以及它的所有祖先总计数量，即为 i 和 j 的公共祖先数量。

对于所有点，我们即可用启发式合并统计自己以及下方有多少个颜色为特定值的点。最终算得：

\sum\limits_{A_i=a}\sum\limits_{A_j=a}deep_{\operatorname{LCA}(i,j)}=\sum\limits_{i=1}^n s_{i,a}^2

答案即为：

s_{1,a}\sum\limits_{A_i=a}deep_i-\sum\limits_{i=1}^n s_{i,a}^2

时间复杂度为 O(n\log n)。

#include<bits/stdc++.h>
#define int long long
using namespace std;

int n,c[200005],dep[200005],sz[200005],gson[200005],ans[200005],P[200005],D[200005],otpt,k,sum,curr_dep,rans;
vector<int>mp[200005];bool isnt_root;

void dfs1(int u,int fa){
    int kk=0;
    sz[u]=1;dep[u]=dep[fa]+1;
    for(int o=0;o<mp[u].size();o++){
        int v=mp[u][o];
        if(v==fa){
            continue;
        }
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>kk){
            gson[u]=v;
            kk=sz[v];
        }
    }
}

void add(int u,int fa){
    otpt-=P[c[u]]*P[c[u]];
    P[c[u]]++;D[c[u]]+=dep[u];
    otpt+=P[c[u]]*P[c[u]];
    for(int o=0;o<mp[u].size();o++){
        int v=mp[u][o];
        if(v!=fa){
            add(v,u);
        }
    }
}

void del(int u,int fa){
    otpt-=P[c[u]]*P[c[u]];
    P[c[u]]--;D[c[u]]-=dep[u];
    otpt+=P[c[u]]*P[c[u]];
    for(int o=0;o<mp[u].size();o++){
        int v=mp[u][o];
        if(v!=fa){
            del(v,u);
        }
    }
}

void dfs2(int u,int fa){
    for(int o=0;o<mp[u].size();o++){
        int v=mp[u][o];
        if(v==fa||v==gson[u]){
            continue;
        }
        dfs2(v,u);
        del(v,u);
    }
    if(gson[u]){
        dfs2(gson[u],u);
    }
    if(u!=1)isnt_root=1;
    curr_dep=dep[u];
    for(int o=0;o<mp[u].size();o++){
        int v=mp[u][o];
        if(v!=fa&&v!=gson[u]){
            add(v,u);
        }
    }
    otpt-=P[c[u]]*P[c[u]];
    P[c[u]]++;D[c[u]]+=dep[u];
    otpt+=P[c[u]]*P[c[u]];
    rans+=otpt;
}

signed main(){
    cin>>n;
    for(int i=1;i<n;i++){
        int a,b;
        cin>>a>>b;
        mp[a].push_back(b);
        mp[b].push_back(a);
    }
    for(int i=1;i<=n;i++){
        cin>>c[i];
    }
    dfs1(1,1);
    dfs2(1,1);
    rans*=-1;
    for(int i=1;i<=n;i++){
        rans+=D[i]*P[i];
    }
    cout<<rans;return 0;
}