Leetcode 3600. Maximize Spanning Tree Stability with Upgrades

• Leetcode 3600. Maximize Spanning Tree Stability with Upgrades
  • 1. 解题思路
  • 2. 代码实现

• 题目链接: 3600. Maximize Spanning Tree Stability with Upgrades

1. 解题思路

这一题核心思路就是一个二分法的思路。我们定义函数is_possible(x)，表示是否存在一个树的构造，使得任意一条边的长度均不少于$x$。

显然，这里有两个前提条件就是：

• 所有$must=1$的边本身不能够成环；
• 至少存在一种构造能够将所有的点连接起来；

如果这两者均不成立，那么直接返回$-1$即可。

然后，我们只需要找到上下确界，然后进行二分查找即可。这个上确界显然就是$must=1$的边的最小值。而下确界显然就是$0$。

于是，问题的核心就变成考察对于任意的$x$，是否存在一种构造使得$k$次升级之内能够获得一个满足条件的树了。而这个我们可以通过DSU算法来求得，只需要将所有满足条件的边全部连起来，看看其是否能将所有的点连通即可。

而关于DSU算法的具体内容，网上反正有很多了，我自己也有一篇水文《经典算法：并查集（DSU）结构简介》作为备忘，有兴趣的读者自己了解下就行了。

2. 代码实现

给出python代码实现如下：

class DSU:
    def __init__(self, N):
        self.node = N
        self.cluster = N
        self.root = [i for i in range(N)]
        
    def find(self, k):
        if self.root[k] != k:
            self.root[k] = self.find(self.root[k])
        return self.root[k]
    
    def union(self, a, b):
        x = self.find(a)
        y = self.find(b)
        if x != y:
            self.root[y] = x
            self.cluster -= 1
        return
    
def copy_dsu(dsu):
    n = dsu.node
    new_dsu = DSU(n)
    new_dsu.root = deepcopy(dsu.root)
    new_dsu.cluster = dsu.cluster
    return new_dsu

class Solution:
    def maxStability(self, n: int, edges: List[List[int]], k: int) -> int:
        must = [(u,v,w) for u, v, w, m in edges if m == 1]
        candi = [(u,v,w) for u, v, w, m in edges if m == 0]
        candi = sorted(candi, key=lambda x: x[2], reverse=True)
        dsu = DSU(n)
        l, r = 0, math.inf
        for u, v, w in must:
            if dsu.find(u) == dsu.find(v):
                return -1
            dsu.union(u, v)
            r = min(r, w)
        r = 2*candi[0][2] if r == math.inf else r
        
        def is_possible(m):
            _dsu = copy_dsu(dsu)
            cnt = 0
            for u, v, w in candi:
                if w * 2 < m or (w < m and cnt >= k):
                    break
                if _dsu.find(u) == _dsu.find(v):
                    continue
                if w < m:
                    cnt += 1
                _dsu.union(u, v)
                if _dsu.cluster == 1:
                    break
            return _dsu.cluster == 1

        if is_possible(r):
            return r
        elif not is_possible(l):
            return -1
        while r-l > 1:
            m = (l+r) // 2
            if is_possible(m):
                l = m
            else:
                r = m
        return l

提交代码评测得到：耗时1271ms，占用内存74.06MB。