Problem Statement in English

You’re given an undirected graph with n nodes and edges where each edge has a stability s and a flag m indicating if it’s mandatory (1) or optional (0). You can upgrade up to k optional edges, doubling their stability.

Your task is to find the maximum possible minimum stability of a spanning tree that can be formed using the edges, considering the upgrades.

Approach

We could reverse the problem and try and guess the minimum stability guess of the spanning tree. Then we can check if it’s possible to form a spanning tree with that minimum stability using the edges and upgrades.

So in each pass we will do the following:

  • Check if all mandatory edges have stability >= guess. If not, return False immediately since we can’t use those edges.

  • Take each mandatory edge and union its endpoints in a Union-Find structure. If we find a cycle among mandatory edges, return False immediately since we can’t form a tree.

  • For optional edges that already have stability >= guess, we can use them for free. Union their endpoints as well - so long as they don’t create a cycle.

  • For optional edges that have stability < guess but can be upgraded to reach guess (i.e. s * 2 >= guess), we can consider upgrading them. We will use them if they connect two different components in the Union-Find structure, and we have upgrades left.

  • Finally, we check if we have exactly 1 component in the Union-Find structure, which means we can form a spanning tree with minimum stability >= guess.

One little speedup we can achieve is by checking if it’s even possible to form a spanning tree with the given edges (without any upgrades) - so a a spanning tree with min stability $0$. We can do this by calling check(0) at the start. If that returns False, we can return -1 immediately.

And we’re done!

Solution in Python

  • Union Find + Binary Search Approach

class UnionFind:
    def __init__(self, n):
        # Using a list is faster and safer for fixed n
        self.parent = list(range(n))
        self.rank = [0] * n
        self.components = n

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        pa = self.find(a)
        pb = self.find(b)
        if pa == pb:
            return False
        if self.rank[pa] > self.rank[pb]:
            self.parent[pb] = pa
        elif self.rank[pb] > self.rank[pa]:
            self.parent[pa] = pb
        else:
            self.parent[pb] = pa
            self.rank[pa] += 1
        self.components -= 1
        return True

class Solution:
    def maxStability(self, n: int, edges: List[List[int]], k: int) -> int:
        mandatory = []
        optional = []
        
        for u, v, s, m in edges:
            if m == 1:
                mandatory.append((u, v, s))
            else:
                optional.append((u, v, s))

        def check(guess):
            # 1. Check if all mandatory edges meet the threshold
            # Mandatory edges CANNOT be upgraded.
            dsu = UnionFind(n)
            for u, v, s in mandatory:
                if s < guess:
                    return False
                if not dsu.union(u, v): # Cycle in mandatory edges
                    return False
            
            upgrades_left = k
            
            # 2. Use optional edges that are ALREADY >= guess (Free)
            for u, v, s in optional:
                if s >= guess:
                    dsu.union(u, v)
            
            # 3. Use optional edges that can REACH guess via upgrade
            for u, v, s in optional:
                if s < guess and s * 2 >= guess and upgrades_left > 0:
                    if dsu.find(u) != dsu.find(v):
                        dsu.union(u, v)
                        upgrades_left -= 1
            
            return dsu.components == 1

        # Binary Search Range
        # Low is 0 (or -1 if impossible), High is max possible strength (2 * max_s)
        low, high = 0, 2000000000 
        ans = -1

        # Initial connectivity check: can we even form a tree?
        # A quick way is to check if 'check(0)' works.
        if not check(0):
            return -1

        while low <= high:
            mid = (low + high) // 2
            if check(mid):
                ans = mid
                low = mid + 1
            else:
                high = mid - 1
        
        return ans

Complexity

  • Time: $O((m + n) \log M)$
    Where $m$ is the number of edges, $n$ is the number of nodes, and $M$ is the maximum possible stability. The $\log M$ factor comes from the binary search.

  • Space: $O(n)$
    Since we are using a Union-Find data structure that requires space proportional to the number of nodes.

Mistakes I Made

I had to loko this up :(

And we are done.