树上并查集

作者: | 更新日期:

一个很神奇的优化,有必要学一下。

本文首发于公众号:天空的代码世界,微信号:tiankonguse

零、背景

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

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

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

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

一、问题

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

最小的未出现的正整数有个专业名称,叫做 mex(Minimum excludant)。

二、暴力方法

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

极端情况下,考虑这个树是一个链,从下到上集合大小依次是 1 到 n。
集合合并 n 次,总复杂度就是 1+2+3+...+nO(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

本文首发于公众号:天空的代码世界,微信号:tiankonguse
如果你想留言,可以在微信里面关注公众号进行留言。

关注公众号,接收最新消息

tiankonguse +
穿越