P2664 树上游戏

知识点：点分治

原题面 Luogu。

简述

给定一棵 \(n\) 个节点的树，节点 \(i\) 的颜色为 \(c_i\)。定义 \(s(i,j)\) 表示节点 \(i\) 到节点 \(j\) 简单路径上的颜色种类数，定义 \(sum_i = \sum\limits_{1\le j\le n} s(i,j)\)，求：\(sum_1\sim sum_n\)。
\(1\le n,c_i\le 10^5\)。
1S，128MB。

分析

套路地考虑点分治，先钦定一个根节点 \(root\)，仅考虑处理过根节点的路径，并递归处理不过根节点的路径。

考虑先处理出以根为端点的链的颜色种类数，并更新 \(sum_{root}\)。从路径的角度统计贡献并不容易，考虑从颜色的角度统计。从 \(root\) 开始 dfs，并维护一个 \(\operatorname{cnt}\) 数组储存当前节点到根路径上各颜色出现的次数。若 dfs 到节点 \(u\) 时有 \(\operatorname{cnt}_{c_u} = 1\)，说明以 \(u\) 的子树节点为一端点，根为一端点的路径上都会出现颜色 \(c_u\)，\(sum_{root}\) 应增加 \(\operatorname{size}_{u}\)。

再考虑对非根节点的贡献。考虑在上一步中预处理一些值：所有链的贡献量 （即\(sum_{root}\) 的增量） \(S\)，各颜色有贡献的路径数量 \(f\)。若在上一步中 dfs 到节点 \(u\) 时有 \(\operatorname{cnt}_{c_u} = 1\)，则令 \(S\gets S+\operatorname{size}_u\)，\(f_{c_u}\gets {f}_{c_u} + \operatorname{size}_u\)。

之后考虑 dfs 处理非根节点的答案。对于 dfs 到的节点 \(u\)，设 \(u\) 所在的 \(root\) 的儿子为 \(son\)，过根节点且过节点 \(u\) 的路径显然有 \(\operatorname{size}_{root} - \operatorname{size}_{son}\) 条。将这样的路径 \(u\rightarrow v\) 拆成链 \(root\rightarrow u\) 与链 \(root \rightarrow v\) 分别考虑它们对 \(sum_u\) 的贡献。

对于链 \(root \rightarrow u\)，它在每一条有贡献的路径中都被包含。考虑在 dfs 枚举 \(u\) 时通过维护当前节点到根路径上各颜色出现的次数，处理出其中的颜色种类数，乘上路径数量 \(\operatorname{size}_{root} - \operatorname{size}_{son}\) 即为其贡献。

对于所有链 \(root \rightarrow v\)，先考虑从所有链的贡献 \(S\) 与各颜色的贡献 \(f\) 中减去子树 \(son\) 的贡献。具体地，考虑枚举子树 \(son\) 的节点 \(x\)。若某颜色在 \(root\rightarrow x\) 的路径上是第一次出现，则令 \(S\gets S-\operatorname{size}_u\)，\(f_{c_u}\gets {f}_{c_u} - \operatorname{size}_u\)。处理完成后 \(S\) 与 \(f\) 仅包含其他子树的信息。之后进行 dfs 枚举节点 \(u\) 并更新 \(sum_u\)，若某颜色在路径 \(root \rightarrow u\) 上出现，说明该颜色的贡献已经在上一步统计 \(root \rightarrow u\) 时已经统计过，的则令 \(S \gets S - f_{c_u}\)。此时所有链 \(root \rightarrow v\) 对 \(sum_u\) 的贡献即为 \(S\)。统计两种链的贡献可以在一个 dfs 中完成。

注意统计完成后需要还原子树 \(son\) 的贡献。

上述过程均可在 \(O(n)\) 的时间复杂度内完成，总时间复杂度是常数巨大的 \(O(n\log n)\)。

代码

//知识点：点分治
/*
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_) { //求得所有链的贡献
  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_) { //更改子树的贡献
  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_) { //统计链 root -> u 与 链 root -> v 的贡献 
  cnt[col[u_]] ++;
  if (cnt[col[u_]] == 1) { //颜色 col[u] 在链 root -> u 中出现过，需要去除该颜色在其他链 root -> v 中的贡献
    sum -= colsz[col[u_]]; 
    ++ num; //num 为链 root -> u 上的颜色种类数
  }
  ans[u_] += sum + num * deltasz; //更新 
  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) { //回溯时还原 
    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; //对 root 的贡献 
  for (int i = head[u_]; i; i = ne[i]) {
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    cnt[col[u_]] ++;  //清除子树 v 的贡献。根 u 为必经点，需要初始化其数量 
    sum -= sz[v_], colsz[col[u_]] -= sz[v_];
    Modify(v_, u_, -1);
    cnt[col[u_]] --;

    deltasz = sz[u_] - sz[v_]; //路径总数 
    Dfs2(v_, u_);

    cnt[col[u_]] ++; //还原子树 v 的贡献 
    sum += sz[v_], colsz[col[u_]] += sz[v_];
    Modify(v_, u_, 1);
    cnt[col[u_]] --;
  }
  sum = num = 0; //注意清空 
  Clear(u_, fa_);

  for (int i = head[u_]; i; i = ne[i]) { //分治 
    int v_ = v[i];
    if (v_ == fa_ || vis[v_]) continue;
    sumsz = sz[v_];
    root = 0, maxsz[root] = kN;
    CalcSize(v_, u_), CalcSize(root, 0), 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), CalcSize(root, 0), Dfs(root, 0);
  for (int i = 1; i <= n; ++ i) printf("%lld\n", ans[i]);
  return 0;
}