树上并查集

零、背景

上次在 Leetcode 258 比赛最后一题的时候，我提到：本想暴力合并集合，但是一计算复杂度，时间复杂度会爆炸。
具体来说，时间复杂度会退化到 O(n^2)。

由于当时所有节点值互不相同，我就想到了树转链表的贪心方法。

后来查了下资料，发现有重复值时，使用暴力方法也是可以做的，不过需要增加一个小优化。

今天分享一个小优化，可以把复杂度降到O(n log(n))，从而可以解决这道题。

一、问题

给一个有根树，每个节点有一个正整数值（可能相同）。
对于每个子树，问子树的节点值里面，最小的未出现的正整数是多少。

最小的未出现的正整数有个专业名称，叫做 mex（Minimum excludant）。

二、暴力方法

对于这个问题，最暴力的方法是递归求出每个儿子的节点集合，然后合并得到当前节点的集合，返回出去。

极端情况下，考虑这个树是一个链，从下到上集合大小依次是 1 到 n。
集合合并 n 次，总复杂度就是 1+2+3+...+n 为 O(n^2)了。

三、理想优化

面对上面的集合合并，我们很容易想到：能不能第一个儿子直接复用当前要返回的集合呢？
这样第一个儿子就不需要复制一遍集合了。

我们假设这个是一个完全二叉树，按照这个假设去计算复杂度，会发现复杂度竟然是 O(n log(n)) 的。

证明也很简单。

由于是完全二叉树，树高是 log(n)，最底层有 n/2 个叶子节点。
倒数第二层有 n/4 个节点，由于有个儿子复用集合，只有一个儿子的集合需要复制，复杂度(n/2) / 2。
同理，倒数第三层有n/8个节点，一半的子孙节点需要复制，次数为 (n/2 + n/4)/2。
同理，倒数第四层有n/16个节点，一半的子孙节点需要复制，次数为(n/2 + n/4 + n/8) /2。
以此递推，到达第一层，需要复制的次数是 n / 2。

我们把所有次数相加，可以得到下面的公式。

f(n) = (n/2) / 2 + (n/2 + n/4)/2 + ...
2 f(n) = n/2 + (n/2 + n/4) + ...
2 f(n) / n = 1/2 + (1/2 + 1/4) + ...
2 f(n) / n = 1-1/2 + 1-1/4 + ...
2 f(n) / n = log(n) - (1/2 + 1/4 + ...)
2 f(n) / n = log(n) - (1 - 1/n)
2 f(n) = n log(n) - n + 1
f(n) <= n log(n)

四、重边优化

在算法领域，有个重边轻边的概念。

对于一个树，到节点数量最多的儿子的边，称为重边，其他边称为轻边。

上一小节已经证明了完全二叉树的优化可以到达 n log(n)的复杂度。

实际上，我们还可以证明，每次复用重边的集合，轻边集合全部复制一份，一样可以到达 n log(n) 的复杂度。

这里就不证明了，网上资料很多，大家可以自己去查找下。

五、树上并查集

有了重边贪心的复杂度理论支持，树上并查集的复杂度就可以由 O(n^2) 降低到 O(n log(n)) 了。

为了找到重边，可以先 DFS 计算每个子树节点的个数。

int DfsCount(int root){
    int num = 1;
    for(auto c: tree[root]) {
        num += DfsCount(c);
    }
    return child_nums[root] = num;
}

为了快速找到重边，可以预处理对儿子进行排序，使得重边就是第一个儿子。

void SortTree(){
    for(int i=0;i<n;i++){
        if(child_nums[i] <= 2) continue; // 包含自己
        vector<int>& child_nums = this->child_nums;
        sort(tree[i].begin(), tree[i].end(), [&child_nums](int a, int b){
            return child_nums[a] > child_nums[b];
        });
    }
}

然后，就可以对重边贪心复用，从而实现树上并查集。

void DfsDsu(int root, set<int>& s){
    int one_ans = 1;
    
    for(auto c: tree[root]) {
        set<int> cs;
        DfsDsu(c, cs);
        
        if(s.empty()) { // 第一个是重边
            s.swap(cs);
        }
        
        for(auto v: cs) { // 合并子树到重边集合
            s.insert(v);
        }
        
        one_ans = max(one_ans, ans[c]);
    }
    
    s.insert(nums[root]); //插入当前子树根
    
    while(s.count(one_ans)) {
        one_ans++;
    }
    ans[root] = one_ans;
}

六、最后

当然，我们在做集合合并的时候，还可以做一个优化：去掉连续前缀。
比如一个集合的 mex 是 x，那么集合中显然存在 [1, x)。
既然我们知道存在连续前缀了，就没必要储存在集合里了，这样还可以降低合并的次数，从而进一步降低复杂度。

你之前听说过树上并查集吗？

加油，算法人。

《完》

-EOF-

本文公众号：天空的代码世界
个人微信号：tiankonguse
公众号ID：tiankonguse-code