Knowledge points: divide and conquer

Original problem surface Luogu.

Brief description

Given a tree with \ (n \) nodes, the color of node \ (I \) is \ (c_i \). Definition \ (s(i,j) \) represents the number of color types on the simple path from node \ (I \) to node \ (j \). Definition \ (sum_i = \ sum \ limits {1 \ Le j \ Le n} s(i,j) \), find: \ (sum_1\sim sum_n \).
\(1\le n,c_i\le 10^5\).
1S，128MB.

analysis

Consider the rule of point division in a routine way. First specify a root node \ (root \), only consider the path that has processed the root node, and recursively process the path that cannot be processed by the root node.

Consider processing the number of color types of the chain with the root as the endpoint first and updating \ (sum {root} \). It is not easy to count the contribution from the perspective of path. Consider counting from the perspective of color. Start dfs from \ (root \), and maintain a \ (\ operatorname{cnt} \) array to store the number of colors on the current node to the root path. If there is \ (\ operatorname{cnt} {c_} = 1 \) when dfs is connected to node \ (U \), it means that the subtree node of \ (U \) is the end point, and the color \ (c_ \) will appear on the path with the root as the end point, and \ (sum {root} \) should be added \ (\ operatorname {size} {u} \).

Then consider the contribution to non root nodes. Consider preprocessing some values in the previous step: the contribution of all chains (i.e. \ (sum {root} \) \ (s \), and the number of paths contributed by each color \ (f \). If there is \ (\ operatorname {CNT} {c_} = 1 \) when dfs gets to node \ (u \) in the previous step, make \ (s \ gets S + \ operatorname {size} _ \), \ (f {c_} \ gets {f} {c_} + \ operatorname {size} _ \).

Then consider the answer that dfs handles non root nodes. For the node \ (U \) from dfs, set the son of \ (root \) where \ (U \) is located as \ (son \). There are obviously \ (\ operatorname {size} {root} - \ operatorname {size} {son} \) paths through the root node and node \ (U \). Split such path \ (u\rightarrow v \) into chain \ (root\rightarrow u \) and chain \ (root \rightarrow v \) and consider their contribution to \ (sum_ \) respectively.

For the chain \ (root \rightarrow u \), it is included in every contributing path. When enumerating \ (U \) in dfs, maintain the occurrence times of each color on the current node to the root path, process the number of color types, and multiply it by the number of paths \ (\ operatorname {size} {root} - \ operatorname {size} {son} \).

For all chains \ (root \rightarrow v \), first consider subtracting the contribution of subtree \ (son \) from the contribution of all chains \ (S \) and the contribution of each color \ (f \). Specifically, consider enumerating the nodes \ (x \) of the subtree \ (son \). If a color appears for the first time on the path of \ (root\rightarrow x \), make \ (S \ gets S - \ operatorname {size}_ \), \ (f {C _} \ gets {f} {C _} - \ operatorname {size}_ \). After processing, \ (S \) and \ (f \) only contain the information of other subtrees. After that, dfs enumerates nodes \ (U \) and updates \ (sum_ \). If a color appears on the path \ (root \rightarrow u \), it indicates that the contribution of the color has been counted in the previous step of statistics \ (root \rightarrow u \), then make \ (S \ gets S - f_{c_} \). At this time, the contribution of all chains \ (root \rightarrow v \) to \ (sum_ \) is \ (S \). The statistical contribution of the two chains can be completed in one dfs.

Note that the contribution of atomic tree \ (son \) needs to be added after the statistics are completed.

The above processes can be completed within the time complexity of \ (O(n) \), and the total time complexity is a constant \ (O(n\log n) \).

code

//Knowledge points: divide and conquer
/*
By:Luckyblock
*/
#include <algorithm>
#include <cctype>
#include <cstdio>
#include <cstring>
#define LL long long
const int kN = 1e5 + 10;
//=============================================================
int n, m, e_num, col[kN], head[kN], v[kN << 1], ne[kN << 1];
int root, sumsz, sz[kN], maxsz[kN];
LL deltasz, num, sum, cnt[kN], colsz[kN], ans[kN];
bool vis[kN];
//=============================================================
inline int read() {
  int f = 1, w = 0;
  char ch = getchar();
  for (; !isdigit(ch); ch = getchar()) if (ch == '-') f = -1;
  for (; isdigit(ch); ch = getchar()) w = (w << 3) + (w << 1) + (ch ^ '0');
  return f * w;
}
void Chkmax(int &fir, int sec) {
  if (sec > fir) fir = sec;
}
void Chkmin(int &fir, int sec) {
  if (sec < fir) fir = sec;
}
void Add(int u_, int v_) {
  v[++ e_num] = v_, ne[e_num] = head[u_], head[u_] = e_num;
}
void CalcSize(int u_, int fa_) {
  sz[u_] = 1, maxsz[u_] = 0;
  for (int i = head[u_]; i; i = ne[i]) {
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    CalcSize(v_, u_);
    Chkmax(maxsz[u_], sz[v_]);
    sz[u_] += sz[v_];
  }
  Chkmax(maxsz[u_], sumsz - sz[u_]);
  if (maxsz[u_] < maxsz[root]) root = u_;
}
void Dfs1(int u_, int fa_) { //Find the contribution of all chains
  sz[u_] = 1;
  cnt[col[u_]] ++;
  for (int i = head[u_]; i; i = ne[i]) {
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    Dfs1(v_, u_);
    sz[u_] += sz[v_];
  }
  if (cnt[col[u_]] == 1) {
    sum += sz[u_];
    colsz[col[u_]] += sz[u_];
  }
  -- cnt[col[u_]];
}
void Modify(int u_, int fa_, int val_) { //Change the contribution of subtree
  cnt[col[u_]] ++;
  if (cnt[col[u_]] == 1) { //
    sum += val_ * sz[u_];
    colsz[col[u_]] += val_ * sz[u_];
  }
  for (int i = head[u_]; i; i = ne[i]) {
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    Modify(v_, u_, val_);
  }
  cnt[col[u_]] --;
}
void Dfs2(int u_, int fa_) { //Count the contribution of chain root - > u and chain root - > v 
  cnt[col[u_]] ++;
  if (cnt[col[u_]] == 1) { //The color col[u] has appeared in the chain root - > u, and the contribution of this color in other chains root - > V needs to be removed
    sum -= colsz[col[u_]]; 
    ++ num; //num is the number of color types on the chain root - > U
  }
  ans[u_] += sum + num * deltasz; //to update 
  for (int i = head[u_]; i; i = ne[i]) {
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    Dfs2(v_, u_);
  }
  if (cnt[col[u_]] == 1) { //Restore on Backtracking 
    sum += colsz[col[u_]];
    -- num;
  }
  -- cnt[col[u_]];
}
void Clear(int u_, int fa_) {
  cnt[col[u_]] = colsz[col[u_]] = 0;
  for (int i = head[u_]; i; i = ne[i]) {
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    Clear(v_, u_);
  }
}
void Dfs(int u_, int fa_) {
  vis[u_] = true;
  Dfs1(u_, fa_);
  ans[u_] += sum; //Contribution to root 
  for (int i = head[u_]; i; i = ne[i]) {
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    cnt[col[u_]] ++;  //Clear the contribution of subtree v. Root u is a required point, and its quantity needs to be initialized 
    sum -= sz[v_], colsz[col[u_]] -= sz[v_];
    Modify(v_, u_, -1);
    cnt[col[u_]] --;

    deltasz = sz[u_] - sz[v_]; //Total number of paths 
    Dfs2(v_, u_);

    cnt[col[u_]] ++; //The contribution of atomic tree v 
    sum += sz[v_], colsz[col[u_]] += sz[v_];
    Modify(v_, u_, 1);
    cnt[col[u_]] --;
  }
  sum = num = 0; //Pay attention to emptying 
  Clear(u_, fa_);

  for (int i = head[u_]; i; i = ne[i]) { //Divide and conquer 
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    sumsz = sz[v_];
    root = 0, maxsz[root] = kN;
    CalcSize(v_, u_), Dfs(root, 0);
  }
}
//=============================================================
int main() { 
  n = read();
  for (int i = 1; i <= n; ++ i) col[i] = read();
  for (int i = 1; i < n; ++ i) {
    int u_ = read(), v_ = read();
    Add(u_, v_), Add(v_, u_);
  }
  
  sumsz = n;
  root = 0, maxsz[root] = kN;
  CalcSize(1, 0), Dfs(root, 0);
  for (int i = 1; i <= n; ++ i) printf("%lld\n", ans[i]);
  return 0;
}