蓝桥杯2023年第十四届省赛A组-颜色平衡树

题目描述

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

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

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

输入格式

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

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

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

输出格式

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

样例输入

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

样例输出

4

提示

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

对于 30% 的评测用例，n ≤ 200，Ci ≤ 200 ；

对于 60% 的评测用例，n ≤ 5000，Ci ≤ 5000 ；

对于所有评测用例，1 ≤ n ≤ 200000，1 ≤ Ci ≤ 200000，0 ≤ Fi < i 。

整体思路

考虑树上启发式合并，求出子树大小 sz⁡、子树颜色出现次数的桶 cnt⁡，以及颜色出现次数的出现次数的桶 ccnt⁡，判断 cnt⁡(u)×ccnt⁡(cnt⁡(u)) 是否等于 sz⁡(u) 即可判断这棵子树是不是颜色平衡树。

#include <bits/stdc++.h>
#define endl '\n'
using namespace std;

const int N = 2e5 + 5;
int c[N], fa[N], sz[N], son[N];
int cnt[N], ccnt[N]; // 子树的颜色出现次数，颜色出现次数的出现次数 
vector<int> e[N];
int ans;

void add(int u, int son)
{
	ccnt[cnt[c[u]]]--;
	cnt[c[u]]++;
	ccnt[cnt[c[u]]]++;
	for(int v : e[u])
	{
		if(v != son)
		{
			add(v, son);
		}
	}
}

void sub(int u)
{
	ccnt[cnt[c[u]]]--;
	cnt[c[u]]--;
	ccnt[cnt[c[u]]]++;
	for(int v : e[u])
	{
		sub(v);
	}
}

void dfs1(int u)
{
	sz[u] = 1;
	for(int v : e[u])
	{
		dfs1(v);
		sz[u] += sz[v];
		if(sz[v] > sz[son[u]])
		{
			son[u] = v;
		}
	}
}

void dfs2(int u, int opt)
{
	for(int v : e[u])
	{
		if(v != son[u])
		{
			dfs2(v, 0);
		}
	}
	if(son[u])
	{
		dfs2(son[u], 1);
	}
	add(u, son[u]);
	if(cnt[c[u]] * ccnt[cnt[c[u]]] == sz[u])
	{
		ans++;
	}
	if(!opt)
	{
		sub(u);
	}
}

int main()
{
	std::ios::sync_with_stdio(false);
	std::cin.tie(nullptr);
	int n;
	cin >> n;
	for(int i = 1; i <= n; i++)
	{
		cin >> c[i] >> fa[i];
		if(fa[i] != 0)
		{
			e[fa[i]].push_back(i);
		}
	}
	dfs1(1);
	dfs2(1, 0);
	cout << ans << endl;
	return 0;
}