【蓝桥杯 2023 省A】颜色平衡树

#树上启发式合并（dsu on tree）

注：树上启发式合并常用于处理统计树上各个结点的信息（比如整合信息之类的）

一、题目

1.题目描述

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

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

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

2.输入格式

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

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

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

3.输出格式

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

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

二、思路

        记录每个结点的颜色c，某种颜色出现的次数cnt，以及该颜色出现次数的出现次数ccnt

        听起来有点绕口，结合定义可能好理解一些：c很好理解，cnt定义是cnt[c[u]]，

        ccnt就是ccnt [ cnt [ c[u] ] ]

1.重链剖分

        即求出每个结点的重儿子son，每个结点u能够从其重儿子son[u]处继承cnt和ccnt的信息

        重儿子就是结点u所有子树中大小最大的儿子，除了重儿子都是轻儿子

2.添加函数add

        定义add（u，dt）表示将u子树的结点以dt的贡献加入到 cnt 和 ccnt 中，其中dt = +1 或 -1。

        实际上add函数就是用来处理轻儿子的，无论是合并信息还是擦出信息，add函数都可以递归处理以该结点（轻儿子）为父节点的整棵子树信息。

        所以时间复杂度关键在于遍历的轻边数量

3.算法流程calc（u，save）

其中u是当前递归到的结点，save是一个是否保存当前贡献的开关：

1. 首先处理所有轻儿子v，递归calc（v，false），也就是递归求出轻儿子子树的答案，并擦出这颗子树的贡献。
2. 如果有重儿子son，递归calc（son[u]，true），也就是递归求出重儿子子树的答案，并保留这棵子树的贡献。
3. 将当前结点u贡献到 cnt 和 ccnt 中。
4. 对于所有的轻儿子v，调用add（v，+1），将轻儿子子树的贡献计入。此时 cnt 和 ccnt 中的信息是u子树的。
5. 统计u子树的答案，判断是否颜色平衡
6. 如果save = false，调用add（v，-1）擦出贡献。

4.时间复杂度

        由于重儿子的定义，根节点到任意一个结点路径的轻边不超过O(nlong n)条。

        一个点会被暴力统计贡献，只有在 calc 搜到这个点，或者搜到这个点的某个作为轻儿子的祖先时才会发生。于是每个点被暴力到的次数为 O(logn)，总复杂度为 O(nlogn)。

三、代码实现

#include<iostream>
#include<vector>
#include<unordered_map>
using namespace std;

const int N = 2e5 + 10;
vector<int> g[N];//邻接表 表示树
int c[N];//每个结点的颜色
int cnt[N], ccnt[N];//每种颜色出现的次数cnt[c[i]]  每种颜色出现次数的出现次数ccnt[cnt[c[i]]]
int sz[N];//每个结点的子树大小
int son[N];//每个结点的重儿子
int res = 0;//结果，满足颜色平衡的子树数量

//dfs，计算每个结点的子树大小，并找到重儿子
void dfs1(int u) {
	sz[u] = 1;//当前结点子树大小初始为1
	for (int v : g[u]) {
		dfs1(v);
		if (sz[v] > sz[son[u]]) son[u] = v;//更新重儿子
		sz[u] += sz[v];//累加子树大小
	}
}

//添加/减少颜色统计cnt ccnt
//实际上add函数就是用来处理轻儿子的，无论是合并信息还是擦出信息，add函数都可以递归处理以该结点（轻儿子）为父节点的整棵子树
void add(int u, int dt) { //dt为+1为添加 dt为-1为减少
	ccnt[cnt[c[u]]]--;
	cnt[c[u]] += dt;
	ccnt[cnt[c[u]]]++;
	for (int v : g[u]) add(v, dt);
}

//calc，统计颜色信息，并判断是否满足平衡条件
void calc(int u, bool keep) {
	//keep指的是当前处理结点u是否是重儿子
	
	//1。先递归处理轻儿子，并清除轻儿子的统计信息
	for (int v : g[u])if (v != son[u]) calc(v, false);

	//2.处理重儿子，并保留重儿子的统计信息
	if (son[u]) calc(son[u], true);

	//3.将当前结点颜色加入统计
	ccnt[cnt[c[u]]]--;
	cnt[c[u]]++;
	ccnt[cnt[c[u]]]++;

	//4.合并轻儿子信息到重儿子中 ，并进行判断颜色平衡
    
    //遍历u结点所有轻儿子，将信息保存到cnt和ccnt中，此时因为处理的是重儿子u，所以信息不会被擦除，即实现了轻儿子信息保存到重儿子中！
	for (int v : g[u]) if(v != son[u]) add(v, 1);
	
    if (ccnt[cnt[c[u]]] * cnt[c[u]] == sz[u]) res++;

	//5.若该节点为轻儿子，清除结点信息
	if (!keep) add(u, -1);
}

int main() {

	int n;
	cin >> n;
	for (int i = 1; i <= n; i++) {
		int fa;
		cin >> c[i] >> fa;
		if (fa != 0) {
			//遍历树时 只存在父节点遍历儿子的情况 儿子无法到达父节点
			g[fa].push_back(i);
		}
	}

	//第一次dfs，计算子树大小 和 重儿子
	dfs1(1);

	//第二次dfs，统计颜色信息，判断平衡条件
	dfs2(1, true);
	
	//输出结果
	cout << res << endl;

	return 0;
}