??? note "例题 3 Luogu P2664 树上游戏" 一棵每个节点都给定颜色的树，定义 $s(i,j)$ 为 $\mathit{i}$ 到 $\mathit{j}$ 的颜色数量，$\mathit{sum_{i}}=\sum_{j=1}^n s(i,j)$。对所有的 $1\leq i\leq n$，求 $sum_i$。（$1 \le n, c_i \le 10^5$）

这道题很考验对点分治思想的理解和应用，适合作为点分治的难度较高的例题和练习题。

首先，我们需要想明白一个转化。题目定义 $\mathit{sum_i}$ 是 $i$ 到所有节点路径上的颜色数量之和，可是如果用这个方法，在点分治中是不好统计答案的，因为这样很难合并从当前根出发的两棵子树的信息。所以我们想到将 $\mathit{sum_i}$ 的意义转化。对于每个颜色 $j$, 其中一个端点为 $i$ 且含有颜色 $j$ 的路径数量记为 $\mathit{cnt_j}$，$\mathit{sum_i}$ 其实就是 $\sum \mathit{cnt_j}$。这一步转化其实就是换了个观察对象，考虑的是每个颜色对 $\mathit{sum_i}$ 的 贡献。而 $\mathit{cnt_j}$ 其实很好处理出来，只需要每遇到一个新颜色，就 $\mathit{cnt_{col_u}}+=\mathit{size_u}$ 即可，其中 $\mathit{size_u}$ 为 u 的子树大小，意味着这个子树里的所有节点都在这个颜色上对 $u$ 的答案有一个贡献。

考虑到点分治过程中，我们只需要分别考虑统计：

1. 子树中以当前根节点为端点的路径对根的贡献
2. lca 为当前根节点的路径对子树内每个点的贡献

1 部分比较好办，由于点分治中，递归层数不超过 $\log{n}$，每一层我们都可以遍历全部子树，这个时候就可以使用 $\mathit{sum_i}$ 的定义式来在遍历子树的过程中顺便统计了。

而针对 2 部分，设当前根节点 $u$ 的一个子节点为 $d$,$d$ 的子树里任取一个点为 $v$，那么 $v$ 的答案可以分为两部分：

1. $(u, v)$ 路径上出现过的颜色，数量设为 $\mathit{num}$，$u$ 除了 $d$ 以外的其他所有子树的总大小设为 $\mathit{siz1}$, 那么这些出现过的颜色对 $v$ 的答案贡献为 $\mathit{num}\times \mathit{siz1}$。
2. $(u, v)$ 路径上没有出现过的颜色 $j$，它们的贡献来自于 $u$ 除了 $d$ 以外的其他所有子树的 $\mathit{cnt_j}$，这部分答案为 $\sum_{j \notin (u, v)} \mathit{cnt_j}$。

以上是全部统计思路，实现细节详见参考代码。

??? note "参考代码"

#include <algorithm>
#include <iostream>
using namespace std;
#define rep(i, a, b) for (int i = (a); i <= (b); ++i)
constexpr int N = 200005;
int h[N], nxt[N * 2], to[N * 2], c[N], gr;

void tu(int x, int y) { to[++gr] = y, nxt[gr] = h[x], h[x] = gr; }

using ll = long long;
int n, nn, siz[N], mn, rt;
bool vis[N];

void get_root(int u, int f) {
  siz[u] = 1;
  int mx = 0;
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (vis[d] || d == f) continue;
    get_root(d, u);
    siz[u] += siz[d];
    mx = max(mx, siz[d]);
  }
  mx = max(mx, nn - siz[u]);
  if (mx < mn) mn = mx, rt = u;
}

ll ans[N], sum;
int cnt[N], v[N];
// sum实时统计的是cnt[i]的和
int nowrt;

void get_dis(int u, int f, int now) {  // now为当前树链上的颜色数量(不含u)
  siz[u] = 1;
  if (!v[c[u]]) {
    sum -= cnt[c[u]];  // 减去在之前子树中已经出现过的颜色信息
    now++;
  }
  v[c[u]]++;
  ans[u] += sum + now * siz[nowrt];  // 统计过u点的路径对u的贡献
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (d == f || vis[d]) continue;
    get_dis(d, u, now);
    siz[u] += siz[d];
  }
  v[c[u]]--;
  if (!v[c[u]]) {
    sum += cnt[c[u]];  // 回溯
  }
}

void get_cnt(int u, int f) {
  if (!v[c[u]]) {
    cnt[c[u]] += siz[u];
    sum += siz[u];  // 将刚遍历过的子树的信息整合到cnt[i]和sum上去
  }
  v[c[u]]++;
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (vis[d] || d == f) continue;
    get_cnt(d, u);
  }
  v[c[u]]--;
}

void clear(int u, int f, int now) {
  if (!v[c[u]]) now++;
  v[c[u]]++;
  ans[u] -= now;
  ans[nowrt] += now;
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (vis[d] || d == f) continue;
    clear(d, u, now);
  }
  v[c[u]]--;
  cnt[c[u]] = 0;
}

void clear2(int u, int f) {
  cnt[c[u]] = 0;
  for (int i = h[u]; i; i = nxt[i]) {
    int d = to[i];
    if (vis[d] || d == f) continue;
    clear2(d, u);
  }
}

int son[N];

void divid(int u) {
  vis[u] = true;
  int tot = 0;
  nowrt = u;
  ans[u]++;
  for (int i = h[u]; i; i = nxt[i]) {
    if (vis[to[i]]) continue;
    son[++tot] = to[i];
  }
  siz[u] = sum = cnt[c[u]] = 1;
  v[c[u]]++;
  rep(i, 1, tot) {  // 统计每个子树和它之前的所有子树中节点组合产生的贡献
    int d = son[i];
    get_dis(d, u, 0);
    get_cnt(d, u);
    siz[u] += siz[d];
    cnt[c[u]] += siz[d];
    sum += siz[d];
  }
  clear2(u, 0);  // 清空数组，记得不可以用memset
  siz[u] = sum = cnt[c[u]] = 1;
  for (int i = tot; i >= 1;
       --i) {  // 统计每个子树和它之后的所有子树中节点组合产生的贡献
    int d = son[i];
    get_dis(d, u, 0);
    get_cnt(d, u);
    siz[u] += siz[d];
    cnt[c[u]] += siz[d];
    sum += siz[d];
  }
  v[c[u]]--;
  clear(u, 0, 0);                      // 清空的同时统计答案
  for (int i = h[u]; i; i = nxt[i]) {  // 继续向下进行点分治
    int d = to[i];
    if (vis[d]) continue;
    nn = siz[d], mn = n + 1, rt = 0;
    get_root(d, u);
    divid(rt);
  }
}

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n;
  int u, v;
  rep(i, 1, n) cin >> c[i];
  rep(i, 2, n) cin >> u >> v, tu(u, v), tu(v, u);
  rt = 0, nn = n, mn = n + 1;
  get_root(1, 0);
  divid(rt);
  rep(i, 1, n) cout << ans[i] << '\n';
  return 0;
}

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=200005;
vector<int>G[N];
int c[N],siz[N],cnt[N];
ll col[N],ans[N],much,sum,num;
int tot,n,rt,rtmaxsiz,all;bool del[N];

void getroot(int x,int xfa){
    siz[x]=1;
    int xmaxsiz=0;
    for(int y:G[x])if(y!=xfa&&!del[y]){
        getroot(y,x);
        siz[x]+=siz[y];
        xmaxsiz=max(xmaxsiz,siz[y]);
    }
    xmaxsiz=max(xmaxsiz,all-siz[x]);
    if(xmaxsiz<rtmaxsiz) rtmaxsiz=xmaxsiz,rt=x;
}
void dfs1(int x,int xfa){
    //要重新dfs一次，不能直接在找根时的树上做（不然子树和之类的会出错）
    //顺便维护各种东西
    siz[x]=1,++cnt[c[x]];
    for(int y:G[x])if(y!=xfa&&!del[y])
        dfs1(y,x),siz[x]+=siz[y];
    if(cnt[c[x]]==1) sum+=siz[x],col[c[x]]+=siz[x];
    --cnt[c[x]];
}
void dfs2(int x,int xfa){
    //把这棵子树里的颜色的影响消除掉
    //顺便更新答案
    ++cnt[c[x]];
    if(cnt[c[x]]==1) sum-=col[c[x]],++num;
    ans[x]+=sum+num*much;
    for(int y:G[x])if(y!=xfa&&!del[y]) dfs2(y,x);
    if(cnt[c[x]]==1) sum+=col[c[x]],--num;
    --cnt[c[x]];
}
void clear(int x,int xfa){
    cnt[c[x]]=col[c[x]]=0;
    for(int y:G[x])if(y!=xfa&&!del[y])
        clear(y,x);
}
void change(int x,int xfa,int val){
    ++cnt[c[x]];
    for(int y:G[x])if(y!=xfa&&!del[y])
        change(y,x,val);
    if(cnt[c[x]]==1) sum+=siz[x]*val,col[c[x]]+=siz[x]*val;
    --cnt[c[x]];
}
void calc(int x){
    //直接带进去乱搞
    dfs1(x,0);
    ans[x]+=sum-col[c[x]]+siz[x];
    for(int y:G[x])if(!del[y]){
        //dfs，然后把各种影响消除掉
        ++cnt[c[x]];
        sum-=siz[y];
        col[c[x]]-=siz[y];
        change(y,x,-1);
        --cnt[c[x]];
        much=siz[x]-siz[y];
        dfs2(y,x);
        ++cnt[c[x]];
        sum+=siz[y];
        col[c[x]]+=siz[y];
        change(y,x,1);
        --cnt[c[x]];
    }
    sum=0,num=0,clear(x,0);
}
void solve(int x){
    calc(x);
    del[x]=true;
    for(int y:G[x])if(!del[y]){
        all=rtmaxsiz=siz[y];getroot(y,x);getroot(rt,x);
        solve(rt);
    }
}
int main(){
    scanf("%d",&n);
    for(int i=1;i<=n;++i) scanf("%d",&c[i]);
    for(int i=1,x,y;i<n;++i){
        scanf("%d%d",&x,&y);
        G[x].push_back(y);
        G[y].push_back(x);
    }
    all=rtmaxsiz=n;getroot(1,0);getroot(rt,0);
    memset(ans,0,sizeof(ans));memset(col,0,sizeof(col));memset(del,0,sizeof(del));
    solve(rt);
    for(int i=1;i<=n;++i) printf("%lld\n",ans[i]);
    return 0;
}