【Luogu P2664】树上游戏

Problem

Description

\(lrb\) 有一棵树，树的每个节点有个颜色。给一个长度为 \(n\) 的颜色序列，定义 \(s(i,j)\) 为 \(i\) 到 \(j\) 的颜色数量。以及

\[sum[i]=\sum_{j=1}^n s(i,j)
\]

现在他想让你求出所有的 \(sum[i]\)

Input Format

第一行为一个整数 \(n\) ，表示树节点的数量

第二行为 \(n\) 个整数，分别表示 \(n\) 个节点的颜色 \(c[1],c[2]……c[n]\)

接下来 \(n-1\) 行，每行为两个整数 \(x,y\)，表示 \(x\) 和 \(y\) 之间有一条边

Output Format

输出 \(n\) 行，第 \(i\) 行为 \(sum[i]\)

Sample

Input

5
1 2 3 2 3
1 2
2 3
2 4
1 5

Output

10
9
11
9
12

Explanation

Explanation for Input

\(sum[1]=s(1,1)+s(1,2)+s(1,3)+s(1,4)+s(1,5)=1+2+3+2+2=10\)

\(sum[2]=s(2,1)+s(2,2)+s(2,3)+s(2,4)+s(2,5)=2+1+2+1+3=9\)

\(sum[3]=s(3,1)+s(3,2)+s(3,3)+s(3,4)+s(3,5)=3+2+1+2+3=11\)

\(sum[4]=s(4,1)+s(4,2)+s(4,3)+s(4,4)+s(4,5)=2+1+2+1+3=9\)

\(sum[5]=s(5,1)+s(5,2)+s(5,3)+s(5,4)+s(5,5)=2+3+3+3+1=12\)

Range

对于 \(40\%\) 的数据，\(n\le 2000\)

对于 \(100\%\) 的数据，\(1\le n,c[i]\le 10^5\)

Algorithm

点分治

Mentality

点分治的神仙题哇天哪，一个个题解看得我那叫一个懵。我还是看神仙的题解才懂的，我这篇题解希望能让您们理解神仙的做法。

首先瞅一眼数据范围 \(10^5......\) 是 \(nlog\) 或 \(nlog^2\) 的标准范围，那么很显然我们不能统计路径，而是应该统计颜色对路径们的贡献。则第一时间发现和树上路径有关，自然上点分。

那么点分怎么搞呢？说来就不简单啊 \(......\) ，对于当前处理的这颗子树，我们记 \(cnt[i]\) 为它的所有子树内到他的路径中包含颜色 \(i\) 的路径条数，那么当处理到一颗子树内的时候，我们统计得到其他子树贡献的所有 \(cnt\) 数组，那么对于该子树内的一个点，先不考虑它到根节点的路径上的颜色，则其它子树内的颜色对当前节点 \(x\) 的贡献肯定为 \(\sum cnt[i]\) (您别告诉我这个看不出来就成) 。

那么这样一来就有问题如下：

• \(cnt\) 数组如何统计。
• 当前节点 \(i\) 到根节点路径上的颜色的贡献如何处理。

我们先解决 \(cnt\) 数组，这个很简单，我们先统计每颗子树的 \(cnt\) 数组，设 \(col[i]\) 为点 \(i\) 的颜色，则每当我们访问到一个之前没有出现过的颜色 \(col[i]\) ，那么 \(cnt[col]+=size[i]\) 。我想这个很好理解，因为子树内的每个节点到根节点的路径上都经过了节点 \(i\) 。所以我们只需要开个 \(book\) 数组记录 \(book[i]\) 为当前节点到根节点路径上颜色为 \(i\) 的节点个数，进入 \(i\) 节点的时候 \(book[col[i]]++\) ，退出时 \(--\) 即可。

如何减去当前处理的子树对 \(cnt\) 数组的贡献呢？当然是在处理这颗子树之前再扫一遍，把贡献去掉，处理完之后又扫一遍，再加回来 = = 。我也觉得很蠢。

第二个问题稍稍难想一点，设当前分治的子树为 \(now\) ，我们先 \(dfs\) 统计贡献，对于当前结点 \(i\) ，我们记录 \(tot=\sum_{col[j]\in col[pre_i]} cnt[col[j]]\) ，也就是将 \(i\) 到根节点路径上的颜色在 \(cnt\) 数组中对应贡献记录下来，由于其他子树内这些颜色的路径到点 \(i\) 也会经过这些点，我们不能重复计算贡献，那就只能牺牲 \(cnt\) 数组的贡献了 \(QwQ\) 。

接着，如果统计的过程中遇到了一个新的颜色，那么 \(tot+=cnt[col[i]]\) ，可以预料到的是，这个新颜色对它的子树内每个节点的贡献肯定是 \(size[root]-size[now]\) ，也即当前分治子树外所有节点都可以通过当前节点到达当前节点子树内的节点，我们带着这个贡献往下 \(dfs\) ，设其为 \(num\) ，在设一个 \(sum=\sum cnt[i]\) 来表示总贡献，那么当我们 \(dfs\) 到点 \(i\) 的时候就变成了如下流程：

• 判断点 \(i\) 的颜色是否出现过，若没有，则 \(tot+=cnt[col[i]],num+=size[root]-size[now]\) 。
• \(ans[i]+=sum-tot+num\)

然后分治递归处理更多的子树就完了！

不理解可以看下代码(码风仙，无空格，不过有注释)

Code

#include <cstdio>
#include <iostream>
using namespace std;
int n, col[100001];
int rt, sum, top, Y, maxs[100001], size[100001], now[100001], cbook[100001],
    cnt[100001];
int head[100001], nx[200001], to[200001];
bool vis[100001], Book[100001];
long long Sum, ans[100001];
void add(int u, int v, int d) {
  to[d] = v, nx[d] = head[u];
  head[u] = d;
}
void getrt(int x, int fa) {
  size[x] = 1, maxs[x] = 0;
  for (int i = head[x]; i; i = nx[i])
    if (to[i] != fa && !vis[to[i]]) {
      getrt(to[i], x);
      size[x] += size[to[i]];
      maxs[x] = max(maxs[x], size[to[i]]);
    }
  maxs[x] = max(maxs[x], sum - size[x]);
  if (maxs[x] < maxs[rt]) rt = x;
}
void getsize(int x, int fa) {
  size[x] = 1;
  for (int i = head[x]; i; i = nx[i])
    if (to[i] != fa && !vis[to[i]]) getsize(to[i], x), size[x] += size[to[i]];
}
void getcol(int x, int fa) {
  if (!now[col[x]]) {
    cnt[col[x]] += size[x];
    Sum += size[x];
  }  //如果没出现过，我们的贡献就 +=size
  if (!Book[col[x]])
    cbook[++top] = col[x],
    Book[col[x]] = true;  //记录一下当前分治的部分总共有哪些颜色
  now[col[x]]++;          //出现次数变更
  for (int i = head[x]; i; i = nx[i])
    if (to[i] != fa && !vis[to[i]]) getcol(to[i], x);
  now[col[x]]--;
}
void delcol(int x, int fa) {
  if (!now[col[x]]) {
    cnt[col[x]] -= size[x];
    Sum -= size[x];
  }
  now[col[x]]++;
  for (int i = head[x]; i; i = nx[i])
    if (to[i] != fa && !vis[to[i]]) delcol(to[i], x);
  now[col[x]]--;
}
void Count(int x, int fa, int num, long long tot) {
  if (!now[col[x]])
    num++, tot += cnt[col[x]];  //如果此颜色首次出现，辣么记录 tot 贡献
  now[col[x]]++;
  ans[x] += Sum - tot + num * Y;  // ans的处理
  for (int i = head[x]; i; i = nx[i])
    if (to[i] != fa && !vis[to[i]]) Count(to[i], x, num, tot);
  now[col[x]]--;
}
void work(int x) {
  getsize(x, 0);  //先把 size 处理出来
  Sum = 0, top = 0;
  getcol(x, 0);  //统计所有子树的 cnt 数组
  for (int i = head[x]; i; i = nx[i])
    if (!vis[to[i]]) {
      now[col[x]]++, delcol(to[i], x), cnt[col[x]] -= size[to[i]],
          Sum -= size[to[i]];  //先减去当前子树的贡献，各种减就对了
      Y = size[x] - size[to[i]],
      Count(to[i], x, 0, 0);  // Y 就是其他子树的节点数
      getcol(to[i], x), now[col[x]]--, cnt[col[x]] += size[to[i]],
          Sum += size[to[i]];  //再把贡献加回来QwQ
    }
  ans[x] += Sum - cnt[col[x]] + size[x];
  for (int i = 1; i <= top; i++)  //将出现过的颜色的贡献统统删掉
    Book[cbook[i]] = false, cnt[cbook[i]] = 0;
}
void solve(int x) {
  vis[x] = true;
  work(x);  //开始计算贡献
  for (int i = head[x]; i; i = nx[i])
    if (!vis[to[i]]) {
      rt = 0, sum = size[to[i]];
      getrt(to[i], x);
      solve(rt);
    }
}
int main() {
  cin >> n;
  for (int i = 1; i <= n; i++) scanf("%d", &col[i]);
  int u, v;
  for (int i = 1; i < n; i++) {
    scanf("%d%d", &u, &v);
    add(u, v, i);
    add(v, u, i + n);
  }
  maxs[rt = 0] = sum = n;
  getrt(1, 0);  //求重心
  solve(rt);
  for (int i = 1; i <= n; i++) printf("%lld\n", ans[i]);
}