蓝桥杯 2023 省A 颜色平衡树

树上启发式合并是一个巧妙的方法。
 

dsu on tree，可以称为树上启发式合并，是一种巧妙的暴力。用一个全局数组存储结果，对于每棵子树，有以下操作：

先遍历轻儿子，处理完轻儿子后将数组清零（要再遍历一次来清零）
遍历重儿子，遍历完不用清零，再遍历，将轻儿子合并到重儿子上去，其合并结果存储于全局数组
用此时的全局数组来计算父亲
重儿子： 对于一个非叶节点u ，设v 是u 的儿子，且以v 为根的子树包含的节点比以u 的其他儿子为根的子树包含的节点都多，则称v 为u 的重儿子

轻儿子： 对于一个非叶节点u ，在u 的各个儿子中，除了重儿子，都是轻儿子

这题还有个关键就是如何判断是否所有存在的颜色个数相同
有个好思路就是，设立全局变量min和max，表示的是子树u的存在的的颜色的最多数量和最小数量，若min==max，则就是满足条件的子树。

 

#include <iostream>
#include <bits/stdc++.h> // 包含常用的标准库头文件
using namespace std;

int n; // 存储节点数量

// 结构体定义边
struct edge
{
    int y, next; // 存储节点编号和下一条边的位置
} e[400002]; // 全局数组，存储边信息

int first[200001]; // 存储每个节点的第一条边的位置
// 第一个参数为节点编号，第二个参数为节点的颜色编号
int col[200001];
// 存储每个节点的子树大小和重儿子
int size[200001], mson[200001];
// 存储每个颜色对应的节点数量，以及该数量对应的颜色数量
int t[2][200010], mi, ma;
int ans; // 存储结果

// 添加边的函数
void add_edge(int x, int y)
{
    static int len = 0; // 静态变量，用于记录边的数量
    e[++len] = {y, first[x]}; // 将边的信息存入数组中
    first[x] = len; // 更新节点的第一条边的位置为当前边的位置
}

// 深度优先搜索，用于计算每个节点的子树大小和重儿子
void dfs0(int x)
{
    size[x] = 1; // 初始化当前节点的子树大小为1
    // 遍历当前节点的每一条边
    for (int i = first[x]; i != 0; i = e[i].next)
    {
        int y = e[i].y; // 获取当前边连接的节点编号
        dfs0(y); // 递归调用，计算子树的大小
        // 更新当前节点的重儿子
        if (size[y] > size[mson[x]])
            mson[x] = y;
        size[x] += size[y]; // 更新当前节点的子树大小
    }
}

// 添加节点颜色数量信息到数组t中
void add(int x)
{
    t[1][t[0][x]]--; // 将原来颜色数量对应的颜色数量减1
    t[0][x]++; // 将当前颜色数量加1
    t[1][t[0][x]]++; // 将新的颜色数量对应的颜色数量加1
    // 更新颜色数量的最小值和最大值
    if (t[0][x] < mi)
        mi = t[0][x];
    if (t[0][x] > ma)
        ma = t[0][x];
    // 如果颜色数量为0，则颜色数量最小值加1
    if (!t[1][mi])
        mi++;
}

// 删除节点颜色数量信息从数组t中
void del(int x)
{
    t[1][t[0][x]]--; // 将原来颜色数量对应的颜色数量减1
    t[0][x]--; // 将当前颜色数量减1
    t[1][t[0][x]]++; // 将新的颜色数量对应的颜色数量加1
    // 更新颜色数量的最小值
    if (t[0][x] && t[0][x] < mi)
        mi = t[0][x];
    // 如果颜色数量为0，则颜色数量最大值减1
    if (!t[1][ma])
        ma--;
}

// 遍历以x为根节点的子树，根据type类型进行颜色数量的操作
void dfs1(int x, int type)
{
    if (type == 0)
        del(col[x]); // 删除节点颜色数量信息
    else
        add(col[x]); // 添加节点颜色数量信息
    // 遍历当前节点的每一个子节点
    for (int i = first[x]; i; i = e[i].next)
    {
        int y = e[i].y; // 获取当前边连接的节点编号
        dfs1(y, type); // 递归调用，处理子节点的信息
    }
}

// 遍历以x为根节点的子树，统计颜色数量相等的子树数量
void dfs2(int x)
{
    // 遍历当前节点的每一个子节点
    for (int i = first[x]; i; i = e[i].next)
    {
        int y = e[i].y; // 获取当前边连接的节点编号
        // 如果当前节点是重儿子，则跳过
        if (y == mson[x])
            continue;
        dfs2(y); // 递归调用，处理非重儿子的子树
        dfs1(y, 0); // 处理非重儿子的节点颜色数量信息
    }
    // 如果存在重儿子，则处理重儿子的子树
    if (mson[x])
        dfs2(mson[x]);
    // 遍历当前节点的每一个子节点
    for (int i = first[x]; i; i = e[i].next)
    {
        int y = e[i].y; // 获取当前边连接的节点编号
        // 如果当前节点是重儿子，则跳过
        if (y == mson[x])
            continue;
        dfs1(y, 1); // 处理非重儿子的节点颜色数量信息
    }
    add(col[x]); // 添加当前节点的颜色数量信息
    // 如果颜色数量最小值等于颜色数量最大值，说明当前子树中的颜色数量相等
    if (mi == ma)
        ans++;
}

int main()
{
    cin >> n; // 输入节点数量
    int y; // 存储父节点编号
    for (int i = 1; i <= n; i++)
    {
        cin >> col[i] >> y; // 输入节点颜色和父节点编号
        add_edge(y, i); // 添加边信息
    }
    dfs0(1); // 计算每个节点的子树大小和重儿子
    mi = 1; // 初始化颜色数量的最小值为1
    dfs2(1); // 遍历树的根节点，并统计颜色数量相等的子树数量
    cout << ans; // 输出结果
    return 0;
}

进一步注释：
 

len配合first数组，可以记录边的序号。

edge 结构体定义了树的边的信息，其中 y 和 next 分别代表了以下含义：

• y：表示当前边所连接的节点的编号。在树的边结构中，每条边连接两个节点，y 就是其中一个节点的编号，表示了这条边连接到了哪个节点。

• next：表示当前节点 x 的下一条邻接边在边数组中的位置。因为这是一个邻接表的表示方式，所以 next 记录了节点 x 的下一条邻接边在边数组中的位置，以便能够快速地遍历节点 x 的所有邻接边。

t[2][200001]：
       第一行 (t[0][...]) 用于记录每种颜色出现的次数。
       第二行 (t[1][...]) 用于记录不同出现次数的颜色数目。

具体地，t[0][x] 表示颜色 x 出现的次数，而 t[1][y] 表示出现 y 次的不同颜色的数目。通过这个数组，可以实现对每种颜色的统计。