【做题笔记】LuoguP2664 树上游戏

LuoguP2664 - 树上游戏

题目链接

做法：点分治

可以发现，每一个点的答案可以分为两个部分：

1. 这个点为端点向下的路径的答案
2. 经过这个点的路径的答案

第一个部分的答案非常好求：

假设遍历到一个点$x$，设这个点的颜色是$c[x]$，那么包含分治中心到点$x$这段子路径的路径一定包含$c[x]$，如果从分治中心遍历到点$x$的过程中第一次遍历到$c[x]$，那么在分治中心的子树中，以分治中心为端点向下的路径中，有$sz[x]$个路径包含$c[x]$，其中$sz[x]$为点$x$为根节点的子树的大小。点$x$的子树的每一个点，都在$c[x]$颜色上为分治中心提供了$1$的贡献。

这样不断向下dfs的过程中，不断记录哪些颜色已经遍历过，然后统计答案贡献。

第二个部分的答案稍微繁琐一点：

经过分治中心的情况，可以分为：

1. 从其中一个子树出来，经过分治中心，跑到另一个子树当中去。
2. 从分治中心的父节点方向下来，经过分治中心，跑到分治中心的子树下去。

不过这两种情况如果细细想来会发现，第二种情况可以通过上面的节点的第一种情况求得，所以这里对于每一个分治中心只考虑第一种情况。

由于我们只考虑分治中心点的子树（也就是当前的分治块），所以我们可以先通过一次dfs统计出整个分治块的点数、包含的各个颜色的数量。

我们依旧遍历到点$x$，但是这次我们时刻更新的是我们当前遍历的点$x$的答案（很难把这个点的贡献算到其他的所有点的答案上去，所以我们反过来思考）。

现在考虑点$x$可以获得的答案贡献。我们可以先处理出其他子树的点数和各个颜色的数量。从分治中心到其他子树上的点可以得到的颜色，从点$x$跑到分治中心在跑到那些点上，也可以得到那些颜色。设其他子树上第$i$个颜色有$cnt[i]$个，那么对于其他子树上每一个存在的颜色，都可以给点$x$提供$cnt[i]$的贡献。

然后考虑点$x$到分治中心这条路径上的颜色，在这条路径上出现过意味着其他子树的所有点到点$x$的路径都会包含这个颜色，所以这个颜色会直接给点$x$提供其他子树的总大小的贡献。有一些颜色已经在其他子树上出现过，此时之前已经算过其他子树上这个颜色给点$x$带来的贡献$cnt[i]$了，所以在加上其他子树的总大小的时候，要把之前加过的$cnt[i]$减去。

剩下的就是纯纯的点分治了。

#include
using namespace std;
typedef long long LL;
const int N = 100005;

LL n, root, cnt_sum, siz, en;
LL front[N], c[N], mxp[N], sz[N], sum[N], vis[N];
LL cnt[N], ccnt[N], pcnt[N], mk[N], tag[N];
//mxp[i]：第i个点的最大子树大小
//sz[i]：以第i个点为根节点的子树的大小
//sum[i]：第i个点的答案
//vis[i]：第i个点是否已经搜过
//cnt[i]：分治中心的所有子树的颜色信息
//ccnt[i]：用来备份cnt数组
//pcnt[i]：用于统计当前子树的颜色信息，然后从cnt中抠掉
//mk[i]：标记当前搜索的子树中有第i个颜色，用于优化时间复杂度
//tag[i]：当前所遍历的点已经算过了第i个颜色的答案

vector<int> col, col2;
//col：存放当前子树有什么颜色，后续遍历颜色时只遍历这里面的颜色以节省时间
//col2：col的备份 

struct Edge {
	LL v, next;
}e[N * 4];

void addEdge(int u, int v) {
	e[++en] = {v, front[u]};
	front[u] = en;
}

//获取分治块的重心 
void get_root(int u, int f, int total) {
	mxp[u] = 0;
	sz[u] = 1;
	for (int i = front[u]; i; i = e[i].next) {
		int v = e[i].v;
		if (vis[v] or v == f) continue;
		get_root(v, u, total);
		sz[u] += sz[v];
		mxp[u] = max(mxp[u], sz[v]);
	}
	mxp[u] = max(mxp[u], total - sz[u]);
	if (!root or mxp[u] < mxp[root]) {
		root = u;
	}
}

//获取当前分治块的各个点的子树大小 
void get_sz(int u, int f) {
	sz[u] = 1;
	for (int i = front[u]; i; i = e[i].next) {
		int v = e[i].v;
		if (v == f or vis[v]) continue;
		get_sz(v, u);
		sz[u] += sz[v];
	}
}

//dfs获取颜色信息，CNT是传进来的数组 
void dfs(int u, int f, LL* CNT) {
	//之前没有遇到过c[u]这个颜色 
	if (!mk[c[u]]) {
		mk[c[u]] = 1;
		col.push_back(c[u]);
	}
	//c[u]这个颜色没有记过贡献 
	if (tag[c[u]] == 0) CNT[c[u]] += sz[u];
	//标记这个颜色已经在整个子树记录过贡献了
	//之后不用再统计这个颜色了 
	++tag[c[u]];
	for (int i = front[u]; i; i = e[i].next) {
		int v = e[i].v;
		if (v == f or vis[v]) continue;
		dfs(v, u, CNT);
	}
	--tag[c[u]];//递归结束，撤销标记 
}

void fix(int u, int f, LL lst) {
	LL tmp = lst;
	//第一次遇到这个颜色
	//如果其他子树中有c[u]，那么cnt[c[u]]为正，把原来加过的减掉
	//其他子树中没有c[u]时，会有cnt[c[u]] = 0 
	if (tag[c[u]] == 0) tmp += siz - cnt[c[u]];
	++tag[c[u]];
	sum[u] += (cnt_sum + tmp); //更新这个点的贡献，cnt_sum是其他子树的那些颜色的总贡献 
	for (int i = front[u]; i; i = e[i].next) {
		int v = e[i].v;
		if (v == f or vis[v]) continue;
		fix(v, u, tmp);
	}
	--tag[c[u]];
}

void calc(int u) {
	cnt_sum = 0;
	get_sz(u, 0); 
	col.clear(); 
	dfs(u, 0, cnt);//考虑点u为端点，直接向下dfs获得答案 
	for (int x: col) {
		mk[x] = 0;
	}//清空标记数组 
	for (int x: col) {
		cnt_sum += cnt[x];
		ccnt[x] = cnt[x];
	}
	col2 = col;
	sum[u] += cnt_sum;
	LL tmp = cnt_sum; 
	for (int i = front[u]; i; i = e[i].next) {
		int v = e[i].v;
		if (vis[v]) continue;
		tag[c[u]] = 1;
		col.clear();
		dfs(v, u, pcnt);//获取当前子树的颜色信息 
		tag[c[u]] = 0;//清空标记 
		for (int x: col) mk[x] = 0;//清空标记 
		cnt[c[u]] -= sz[v];//减去点u的颜色贡献 
		cnt_sum -= sz[v];//减去点u的颜色贡献 
		for (int x: col) {
			cnt[x] -= pcnt[x];//减去这个子树内的颜色贡献 
			cnt_sum -= pcnt[x];//减去这个子树内的颜色贡献 
		}
		siz = sz[u] - sz[v];
		fix(v, u, 0);//统计分治中心到点u的路径的颜色答案 
		//恢复信息 
		cnt[c[u]] += sz[v];
		cnt_sum = tmp;
		for (int x: col) {
			cnt[x] = ccnt[x];
			pcnt[x] = 0;
		}
	}
	//清空数组 
	for (int x: col2) cnt[x] = 0; 
	col2.clear();
	//标记为已搜索 
	vis[u] = 1;
}

void get_ans(int u) {
	calc(u);
	for (int i = front[u]; i; i = e[i].next) {
		int v = e[i].v;
		if (vis[v]) continue;
		root = 0;
		get_root(v, 0, sz[v]);//找重心 
		get_ans(root);//分治搜索 
	}
}

void main2() {
	cin >> n;
	en = 0;
	for (int i = 1; i <= n; ++i) {
		front[i] = mk[i] = vis[i] = tag[i] = sum[i] = 0;
		cin >> c[i];
	}
	for (int i = 1; i < n; ++i) {
		int x, y;
		cin >> x >> y;
		addEdge(x, y);
		addEdge(y, x);
	}
	mxp[0] = n;
	root = 0;
	get_root(1, 0, n);
	get_ans(root);
	for (int i = 1; i <= n; ++i) {
		cout << sum[i] << '\n';
	}
}

int main() {
	ios::sync_with_stdio(false);
	cin.tie(0); cout.tie(0);
	LL _ = 1;
//	cin >> _;
	while (_--) main2();
	return 0;
}