颜色平衡树

[蓝桥杯 2023 省 A] 颜色平衡树

题目描述

给定一棵树，结点由 $1$ 至 $n$ 编号，其中结点 $1$ 是树根。树的每个点有一个颜色 $C_i$。

如果一棵树中存在的每种颜色的结点个数都相同，则我们称它是一棵颜色平衡树。

求出这棵树中有多少个子树是颜色平衡树。

输入格式

输入的第一行包含一个整数 $n$，表示树的结点数。

接下来 $n$ 行，每行包含两个整数 $C_i,F_i$，用一个空格分隔，表示第 $i$ 个结点的颜色和父亲结点编号。

特别地，输入数据保证 $F_1$ 为 $0$，也即 $1$ 号点没有父亲结点。保证输入数据是一棵树。

输出格式

输出一行包含一个整数表示答案。

样例 #1

样例输入 #1

6
2 0
2 1
1 2
3 3
3 4
1 4

样例输出 #1

4

提示

【样例说明】

编号为 $1,3,5,6$ 的 $4$ 个结点对应的子树为颜色平衡树。

【评测用例规模与约定】

对于 $30\%$ 的评测用例，$n\le 200$，$C_i\le 200$；

对于 $60\%$ 的评测用例，$n\le 5000$，$C_i\le 5000$；

对于所有评测用例，$1\le n\le 200000$，$1\le C_i\le 200000$，$0\le F_i<i$。

完全板子题，没思维含量，没学过根本想不到，无聊至极的题, 建议暴力骗60分
板子题链接 https://codeforces.com/problemset/problem/600/E
做法:DSU ON TREE

#include <iostream>
#include <cstring>
#include <vector>
using namespace std;
using ll = long long;

const int maxn = 2e5 + 10, maxm = maxn * 2;
int n;
vector<int> g[maxn];
int col[maxn], cnt[maxn], siz[maxn], son[maxn];
ll res, sum, mx, totcol;

void dfs1(int u, int fa) {
	siz[u] = 1;
	for (auto j : g[u]) {
		if (j == fa) continue;
		dfs1(j, u);
		siz[u] += siz[j];
		if (siz[j] > siz[son[u]]) son[u] = j;
	}
}

void update(int u, int fa, int sign, int pson) {
	int c = col[u];
	cnt[c] += sign;
	if (sign == 1 && cnt[c] == 1) totcol++;
	if (sign == -1 && cnt[c] == 0) totcol--;
	if (cnt[c] > mx) mx = cnt[c], sum = 1;
	else if (cnt[c] == mx) sum += 1;

	for (auto j : g[u]) {
		if (j == fa || j == pson) continue;
		update(j, u, sign, pson);
	}
}

void dfs2(int u, int fa, int op) {
	for (auto j : g[u]) {
		if (j == fa || j == son[u]) continue;
		dfs2(j, u, 0);
	}
	if (son[u]) dfs2(son[u], u, 1);
	update(u, fa, 1, son[u]);
	if (sum == totcol) res++;
	if (!op) update(u, fa, -1, 0), sum = mx = totcol = 0;
}

int main() {
	cin >> n;
	for (int i = 1; i <= n; i++) {
		cin >> col[i];
		int fa;
		cin >> fa;
		g[fa].push_back(i);
	}
	dfs1(1, -1);
	dfs2(1, -1, 1);
	cout << res << endl;
	return 0;
}