14届蓝桥杯省赛E题——颜色平衡树

一、题目描述

二、题目分析

设颜色平衡树的节点有n个，颜色种类为p种，每种颜色的出现次数均为q,则n=p*q。换句话说，如果一棵树的出现次数最多的颜色们的出现次数之和等于该树的节点个数，那么这棵树是颜色平衡树。

为了降低遍历次数，采用树上启发式合并，定义树中节点最多的子树为重子树，其余为轻子树。

1. 遍历轻子树，得到轻子树中的平衡树，但不保留颜色记录；
2. 遍历重子树，得到重子树中的平衡树，保留颜色记录；
3. 再次遍历轻子树，将轻子树节点的颜色记录贡献到重子树的颜色记录中；
4. 判断这树是否为颜色平衡树；

三、代码

#include<bits/stdc++.h>
using namespace std;
#define N 200050

int color[N];//节点n的颜色
int heavy_son[N];//节点n的重儿子
int Size[N];//节点n的大小
int num[N];//颜色C出现的次数
int max_color_num;//出现次数最多的 颜色 出现的次数
int sum;//出现次数最多的 颜色们 的总出现次数
int ans;//颜色平衡树的数量
vector<int> child[N];//节点n的孩子

void Init_dfs(int node)
//计算每个子树的大小，记录在size[node]中；并且找到每个子树的重子树,记录在heavy[node]中。
{
	Size[node] = 1;
	for (auto it : child[node])
	{
		Init_dfs(it);
		Size[node] += Size[it];
		if (Size[it] > Size[heavy_son[node]])
			heavy_son[node] = it;
	}
}

void add(int node, bool option)
{
	int this_color = color[node];
	num[this_color]++;
	if (num[this_color] > max_color_num)
	{
		max_color_num = num[this_color];
		sum = num[this_color];
	}
	else if (num[this_color] == max_color_num)
		sum += num[this_color];

	
	for (auto it : child[node])
	{
		//首次调用时跳过重子树
		if (it == heavy_son[node] && option)
			continue;
		add(it, false);
	}
}

void sub(int node)
//sub函数将递归地从num[200010]中删除一个节点及其所有儿子的信息
{
	int this_color = color[node];
	num[this_color]--;
	for (auto it : child[node])
		sub(it);
}

void dfs(int node, bool option)
{
	for (auto it : child[node])
	//跳过重子树，遍历轻子树，计算轻子树中的结果，但不保留记录
	{
		if (it == heavy_son[node])
			continue;
		dfs(it, false);
	}

	if (heavy_son[node])
	//如果重子树存在，遍历重子树，计算出重子树中的结果，保留结果
		dfs(heavy_son[node], true);

	//遍历轻子树，保留记录，并与上一步相加
	add(node, true);
	//如果出现次数最多的颜色们的出现次数总和等于该树的大小
	if (sum == Size[node])
		ans++;
	//如果不保留该次次数，则对记录撤销
	if (!option)
	{
		sub(node);
		sum = 0;
		max_color_num = 0;
	}
}


int main()
{
	int n; cin >> n;
	for (int i = 1; i <= n; i++)
	{
		int p;
		cin >> color[i] >> p;
		child[p].emplace_back(i);
	}
	Init_dfs(1);
	dfs(1, true);
	cout << ans;
    return 0;
}