树上启发式合并

封面图：https://www.bilibili.com/opus/977994994306514949

引言

树上启发式合并的算法本身并不难，更值得关注的是它的时间复杂度。启发式算法的时间复杂度通常比较玄学，然而在树上启发式合并中，我们可以看到启发式算法是如何极大优化时间复杂度的。本文将通过一道蓝桥杯的题目来讲解树上启发式合并的算法，并分析它的时间复杂度。

例题

洛谷P9233

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

题目描述

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

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

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

输入格式

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

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

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

输出格式

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

样例 #1

样例输入 #1

样例输出 #1

1

4

提示

【样例说明】

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

【评测用例规模与约定】

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

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

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

题解

注意到该题是一道数据离线的树上数颜色的问题。首先可以想到用 dfs ，枚举每一颗子树，分别进行dfs统计颜色数量，然后判断是否是颜色平衡树。但是这样的时间复杂度是 $O(n^2)$ 的，显然不可接受。

注意到在这种算法中，有些点被重复计算了很多次。能不能减少点被计算的次数呢？很容易想到对于一颗子树，如果它的儿子的子树已经统计过了，那么它只需要直接合并儿子的统计结果就可以了。我们可以用一个桶数组来存储每个颜色出现的次数。颜色数量很多，不能为每个点开一个桶，也不能使用 memset 和 memcpy 来初始化和合并桶数组。很容易想到为子节点和父节点分别开一个桶，对于每颗子树，扫描三次，第一次统计子节点的颜色数量到子节点桶，第二次合并子节点的颜色数量到父节点，第三次清除子节点桶以供兄弟节点使用。每次一层统计完时，就交换子节点桶和父节点桶。这样子未免有点麻烦，其实可以只使用一个数组。每个子树依然扫描三次，第一次统计子节点的颜色数量到数组，第二次清除数组以供兄弟节点使用。在所有兄弟节点用完一遍之后，再统计一次子节点的颜色数量到数组。这样就能统计所有子树的颜色数量给到父节点了。

在实现以上算法时，很容易想到的一个优化是在最后一个兄弟节点使用数组时，可以不用清除数组。直接将其它兄弟节点的颜色数量合并到数组上，这样就能免除这个兄弟节点的清除和重新加入数组两边遍历操作。这个优化的效果很大程度上取决于这个“最后一个兄弟节点”的选择。显然，如果我们能够选择一个“最后一个兄弟节点”，使得它的子节点数量最多，那么这个优化的效果就会最大。这就是启发式合并的核心思想。

时间复杂度思考

在这个算法中，一个节点仍然会被遍历多次。在一层中，一个轻节点会被遍历 3 次，一个重节点会被遍历 1 次。轻节点的儿子也会因为轻节点的需要而被遍历多次。但是，重节点却没有这样的问题。当我们分析时间复杂度时，我们可以从每个节点的角度来分析。对于每一个节点，它被遍历到的次数最多是 $O(\log n)$ 的。我们可以从它的轻节点祖先数量来思考。一个节点如果是轻节点，那么它的子树的节点数量一定小于它的父节点的子树的节点数量的一半（如果超过一半，那么它就是重节点了）。由此，从根节点到某个节点，每经过一条轻边，节点的子树的节点数量就会减少一半。所以，一个节点的轻节点祖先数量最多是 $O(\log n)$ 的（否则总节点数就要超过 $n$ 了）。所以整体的时间复杂度是 $O(n\log n)$ 的。

从以上思考可以看出，树上启发式合并确保了每个节点的遍历次数是 $O(\log n)$ 的，这是由树的性质和一个简单的启发式操作决定的。如果不使用启发式合并，那么最糟糕的情况下，时间复杂度仍然会接近 $O(n^2)$。

代码实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65


#include<iostream>
#include<vector>
using namespace std;
const int MAXN=200005,MAXC=200005;
int n,ans=0;
int color[MAXN],cnt[MAXC],ccnt[MAXN],ns[MAXN];
vector<int>son[MAXN];
int chson(int node);
void dfs(int node);
void add(int node,int delta);
int chson(int node)
{
    int result=1;
    for (auto &i:son[node])
    {
        ns[i]=chson(i);
        result+=ns[i];
        if (ns[i]>ns[son[node][0]])
        {
            swap(i,son[node][0]);
        }
    }
    return result;
}
void dfs(int node)
{
    for (auto i=son[node].begin()+1;i<son[node].end();i++)
    {
        dfs(*i);
        add(*i,-1);
    }
    if (son[node].size()>0) dfs(son[node][0]);
    ccnt[cnt[color[node]]]--;
    cnt[color[node]]++;
    ccnt[cnt[color[node]]]++;
    for (auto i=son[node].begin()+1;i<son[node].end();i++)
    {
        add(*i,1);
    }
    if (cnt[color[node]]*ccnt[cnt[color[node]]]==ns[node]) ans++;
}
void add(int node,int delta)
{
    ccnt[cnt[color[node]]]--;
    cnt[color[node]]+=delta;
    ccnt[cnt[color[node]]]++;
    for (auto i:son[node])
    {
        add(i,delta);
    }
}
int main()
{
    cin>>n;
    for (int i=1;i<=n;i++)
    {
        int c,f;
        cin>>c>>f;
        color[i]=c;
        if (f!=0) son[f].push_back(i);
    }
    ns[1]=chson(1);
    dfs(1);
    cout<<ans;
}