LeetCode 3600 - Maximize Spanning Tree Stability with Upgrades

LeetCode Problem 3600

Difficulty: 🔴 Hard
Topics: Binary Search, Greedy, Union-Find, Graph Theory, Minimum Spanning Tree

Solution

Problem Understanding

This problem asks us to determine the maximum possible stability of a spanning tree built from a weighted, undirected graph with certain constraints. Each edge in the graph has a strength and a mandatory flag. A mandatory edge (musti == 1) must appear in the spanning tree, and cannot be upgraded. Optional edges (musti == 0) may be included and can be upgraded at most once. Upgrading doubles the strength of the edge, but there is a limit of k upgrades.

The stability of a spanning tree is defined as the minimum edge strength in the tree. So, our goal is to select edges, optionally upgrading some, to maximize the minimum strength in the spanning tree.

Inputs:

• n: number of nodes (2 ≤ n ≤ 10^5)
• edges: list of [u, v, s, must] (1 ≤ s ≤ 10^5)
• k: maximum number of upgrades (0 ≤ k ≤ n)

Outputs:

• Maximum possible stability, or -1 if a spanning tree cannot be formed (e.g., mandatory edges form a cycle or graph is disconnected).

Constraints imply that brute-force enumeration of all spanning trees is infeasible due to n being up to 10^5. This necessitates a more algorithmic approach using greedy methods, binary search, and Union-Find for efficiency.

Key edge cases:

• All edges are mandatory but form a cycle → no valid spanning tree.
• Graph is disconnected → return -1.
• No upgrades (k = 0) → solution must rely on existing edge strengths.
• Maximum k upgrades are insufficient to make an edge strong enough to improve stability.

Approaches

The brute-force approach would generate all spanning trees, compute their minimum edge strengths, and apply upgrades to optional edges in all combinations to maximize stability. This is correct in principle but impractical: the number of spanning trees grows exponentially with n, and considering upgrade combinations adds a factor of 2^m for optional edges.

The optimal approach leverages two insights. First, the problem can be reframed as checking whether a candidate stability value x is achievable, i.e., can we form a spanning tree where all edges have strength ≥ x (after optional upgrades)? Second, this feasibility check can be combined with binary search over possible stability values. Within each check, we use a Union-Find structure to attempt building a spanning tree, prioritizing mandatory edges and optional edges that meet the threshold (or can be upgraded within the k limit). This transforms the problem into a decision problem that can be solved efficiently.

Algorithm Walkthrough

1. Sort edges: Separate edges into mandatory and optional. For optional edges, keep track of their original strength and potential upgraded strength.
2. Define binary search bounds: The minimum possible stability is 1, and the maximum is the largest possible doubled edge strength (max(si * 2)).
3. Binary search: For each candidate stability mid, check if it is possible to construct a spanning tree with that minimum strength.
4. Feasibility check:

• Initialize Union-Find with n nodes.
• Include all mandatory edges; if they form a cycle, return false.
• Count how many optional edges can already meet mid without upgrades.
• For optional edges below mid but could reach mid if upgraded, select up to k upgrades to make them eligible.
• If, after including mandatory and eligible optional edges, we can connect all nodes (Union-Find has single component) and the number of upgrades does not exceed k, mid is feasible.

5. Update binary search: If feasible, move the lower bound up; otherwise, move the upper bound down.
6. Result: The maximum feasible mid is the answer. If no feasible stability exists, return -1.

Why it works: Binary search allows us to efficiently identify the highest achievable stability. Union-Find guarantees connectivity without cycles. Mandatory edges are always included, and upgrade decisions are made greedily only when needed to meet candidate thresholds.

Python Solution

from typing import List

class UnionFind:
    def __init__(self, n: int):
        self.parent = list(range(n))
        self.rank = [0] * n

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

    def union(self, x: int, y: int) -> bool:
        px, py = self.find(x), self.find(y)
        if px == py:
            return False
        if self.rank[px] < self.rank[py]:
            self.parent[px] = py
        else:
            self.parent[py] = px
            if self.rank[px] == self.rank[py]:
                self.rank[px] += 1
        return True

class Solution:
    def maxStability(self, n: int, edges: List[List[int]], k: int) -> int:
        mandatory = []
        optional = []
        max_strength = 0
        for u, v, s, must in edges:
            max_strength = max(max_strength, s * 2)
            if must:
                mandatory.append((u, v, s))
            else:
                optional.append((u, v, s))

        def can_achieve(stab: int) -> bool:
            uf = UnionFind(n)
            upgrades_needed = 0

            for u, v, s in mandatory:
                if not uf.union(u, v):
                    return False

            eligible_edges = []
            upgrade_candidates = []
            for u, v, s in optional:
                if s >= stab:
                    eligible_edges.append((u, v))
                elif s * 2 >= stab:
                    upgrade_candidates.append((u, v))
            
            for u, v in eligible_edges:
                uf.union(u, v)
            upgrades_needed = min(k, len(upgrade_candidates))
            for i in range(upgrades_needed):
                uf.union(upgrade_candidates[i][0], upgrade_candidates[i][1])
            
            root = uf.find(0)
            return all(uf.find(i) == root for i in range(n))

        left, right = 1, max_strength
        result = -1
        while left <= right:
            mid = (left + right) // 2
            if can_achieve(mid):
                result = mid
                left = mid + 1
            else:
                right = mid - 1

        return result

The Python solution first separates edges, calculates the maximum potential strength, and defines a helper can_achieve to check feasibility using Union-Find. Optional edges are split into those that already meet the threshold and those requiring upgrades. Binary search iterates to find the maximum stability.

Go Solution

package main

func maxStability(n int, edges [][]int, k int) int {
	type edge struct {
		u, v, s, must int
	}
	mandatory := []edge{}
	optional := []edge{}
	maxStrength := 0
	for _, e := range edges {
		u, v, s, must := e[0], e[1], e[2], e[3]
		if s*2 > maxStrength {
			maxStrength = s * 2
		}
		if must == 1 {
			mandatory = append(mandatory, edge{u, v, s, must})
		} else {
			optional = append(optional, edge{u, v, s, must})
		}
	}

	parent := make([]int, n)
	rank := make([]int, n)
	for i := 0; i < n; i++ {
		parent[i] = i
	}

	var find func(int) int
	find = func(x int) int {
		if parent[x] != x {
			parent[x] = find(parent[x])
		}
		return parent[x]
	}

	union := func(x, y int) bool {
		px, py := find(x), find(y)
		if px == py {
			return false
		}
		if rank[px] < rank[py] {
			parent[px] = py
		} else {
			parent[py] = px
			if rank[px] == rank[py] {
				rank[px]++
			}
		}
		return true
	}

	canAchieve := func(stab int) bool {
		for i := 0; i < n; i++ {
			parent[i] = i
			rank[i] = 0
		}
		upgradesNeeded := 0

		for _, e := range mandatory {
			if !union(e.u, e.v) {
				return false
			}
		}

		eligible := [][]int{}
		upgradeCandidates := [][]int{}
		for _, e := range optional {
			if e.s >= stab {
				eligible = append(eligible, []int{e