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;
}