「题解」洛谷 P2664 树上游戏

描述

题目传送门

一棵树，树的每个节点有个颜色。定义 \(\text{s}(i,j)\) 为 \(i\) 到 \(j\) 的颜色数量。求 \(i\in [1,n]\)，\(\sum_{j=1}^n \text{s}(i, j)\)。

思路

转化题意，从 \(u\) 到所有点的路径上经过的颜色数量，等价于对于每种颜色 \(c\)，从 \(u\) 到 \(v\) 的路径经过第 \(c\) 种颜色的 \(v\) 的个数，等价于 \(\sum_{c\in \text{Color}} n\) \(−\) 从 \(u\) 到 \(v\) 的路径不经过第 \(c\) 种颜色的 \(v\) 的个数。

考虑暴力枚举颜色 \(c\)，将树上所有颜色为 \(c\) 的点删去，形成若干个连通块。可以发现，同一个连通块上的两点路径中一定不经过颜色 \(c\)，而不同连通块上的两点路径中一定经过颜色 \(c\)。这时，第 \(c\) 种颜色对每一点的贡献就是其所在连通块的大小。

优化暴力，从根结点开始 DFS，记 DFS 到 \(x\) 结点。维护 \(\text{siz}_x\) 表示以 \(x\) 结点为根的子树大小，\(\text{colcnt}_c\) 表示 DFS 的过程中以 \(c\) 颜色的结点为根的极大子树的大小之和，\(\text{colall}_c\) 存储 DFS 的过程中所有以 \(c\) 颜色的结点为根的极大子树的根结点。「极大同颜色子树」是指某一以 \(c\) 颜色的结点为根的子树，该子树根结点的祖先中，没有颜色为 \(c\) 的点。

遍历 \(x\) 结点的每棵子树 \(y\)，如下图所示，第 \(c\) 种颜色会对红色部分产生红色部分的大小的贡献，即 \(s=\text{siz}_u-\sum \text{siz}_z\)，\(z\) 是图中小三角部分，也是以 \(z\) 结点为根的子树中的极大同颜色子树。

对其树上差分，记 \(\text{cha}\) 为差分数组，将 \(\text{cha}_u+s\)，将 \(\text{cha}_z-s\)，如下图所示。

最后将差分数组还原，结点 \(u\) 的答案就是 \(m\cdot n-\text{cha}_u\)。

考虑红色部分的大小如何求得。在进入子树 \(y\) 之前，记录下 \(\text{colcnt}_c\)，DFS 进入子树。退出后记录现在的 \(\text{colcnt}_c\)，两次的差就是在子树 \(y\) 中小三角部分的大小，红色部分的大小就是子树 \(y\) 的大小 \(-\) 小三角部分的大小。

代码

#include <bits/stdc++.h>
using namespace std;
#define MAXN 100010
int n, tim, c[MAXN], d[MAXN], cha[MAXN], siz[MAXN], dfn[MAXN], colcnt[MAXN], dfnlft[MAXN], dfnrig[MAXN];
vector<int> g[MAXN], colall[MAXN];
int dfs(int u, int fa) {
	siz[u] = 1; dfn[u] = dfnlft[u] = dfnrig[u] = ++tim;
	for (int i = 0; i < g[u].size(); i++) {
		int v = g[u][i];
		if (v == fa) continue;
		int tmp = colcnt[c[u]];
		dfs(v, u);
		dfnlft[u] = min(dfnlft[u], dfnlft[v]); dfnrig[u] = max(dfnrig[u], dfnrig[v]);
		siz[u] += siz[v];
		int mis = colcnt[c[u]] - tmp, add = siz[v] - mis;
		cha[v] += add;
		colcnt[c[u]] += add;
		if (colall[c[u]].empty()) continue;
		for (int ed = colall[c[u]].back(); colall[c[u]].size(); ed = colall[c[u]].back()) {
			if (dfn[ed] < dfnlft[v] || dfn[ed] > dfnrig[v]) break;
			cha[ed] -= add;
			colall[c[u]].pop_back();
		}
	}
	colcnt[c[u]]++;
	colall[c[u]].push_back(u);
}
void get(int u, int fa) {
	for (int i = 0; i < g[u].size(); i++) {
		int v = g[u][i];
		if (v == fa) continue;
		cha[v] += cha[u];
		get(v, u);
	}
}
signed main() {
	cin >> n;
	int sum = 0;
	for (int i = 1; i <= n; i++) {
		cin >> c[i];
		if (!d[c[i]]) sum++;
		d[c[i]] = 1;
	}
	for (int i = 1; i < n; i++) {
		int u, v;
		cin >> u >> v;
		g[u].push_back(v), g[v].push_back(u);
	} 
	dfs(1, 0);
	for (int i = 1; i <= 100000; i++) {
		if (!d[i]) continue;
		int add = n - colcnt[i];
		cha[1] += add;
		if (colall[i].empty()) continue;
		for (int ed = colall[i].back(); colall[i].size(); ed = colall[i].back()) {
			cha[ed] -= add;
			colall[i].pop_back();
		}
	}
	get(1, 0);
	for (int i = 1; i <= n; i++) cout << sum * n - cha[i] << endl;
	return 0;
} 

参考

• P2664 树上游戏 题解 —— b6e0 的洛谷博客