Maximize spanning tree stability with upgrades

Binary Search Greedy Union-Find Graph Theory Minimum Spanning Tree

39.6% acceptance

Feb 25, 2026

You are given an integer n, representing n nodes numbered from 0 to n - 1 and a list of edges, where edges[i] = [ui, vi, si, musti]: ui and vi indicates an undirected edge between nodes ui and vi. si is the strength of the edge. musti is an integer (0 or 1). If musti == 1, the edge must be included in the spanning tree. These edges cannot be upgraded. You are also given an integer k, the maximum number of upgrades you can perform. Each upgrade doubles the strength of an edge, and each eligible edge (with musti == 0) can be upgraded at most once. The stability of a spanning tree is defined as the minimum strength score among all edges included in it. Return the maximum possible stability of any valid spanning tree. If it is impossible to connect all nodes, return -1. Note: A spanning tree of a graph with n nodes is a subset of the edges that connects all nodes together (i.e. the graph is connected) without forming any cycles, and uses exactly n - 1 edges.

Solution

Rust

Time O(n log n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn max_stability(n: i32, edges: Vec<Vec<i32>>, k: i32) -> i32 {
    let n = n as usize;
    let max_s = edges.iter().map(|e| e[2]).max().unwrap_or(0);

    struct Dsu {
      parent: Vec<usize>,
      rank: Vec<usize>,
      components: usize,
    }

    impl Dsu {
      fn new(n: usize) -> Self {
        Dsu {
          parent: (0..n).collect(),
          rank: vec![0; n],
          components: n,
        }
      }

      fn find(&mut self, x: usize) -> usize {
        if self.parent[x] != x {
          self.parent[x] = self.find(self.parent[x]);
        }
        self.parent[x]
      }

      fn union(&mut self, x: usize, y: usize) -> bool {
        let px = self.find(x);
        let py = self.find(y);
        if px == py {
          return false;
        }
        if self.rank[px] < self.rank[py] {
          self.parent[px] = py;
        } else if self.rank[px] > self.rank[py] {
          self.parent[py] = px;
        } else {
          self.parent[py] = px;
          self.rank[px] += 1;
        }
        self.components -= 1;
        true
      }
    }

    let check = |mid: i32| -> bool {
      let mut dsu = Dsu::new(n);

      // Must-include edges: must have s >= mid (cannot be upgraded)
      for e in &edges {
        if e[3] == 1 {
          if e[2] < mid {
            return false;
          }
          if !dsu.union(e[0] as usize, e[1] as usize) {
            return false; // cycle among mandatory edges
          }
        }
      }

      // Free optional edges (no upgrade needed, s >= mid)
      for e in &edges {
        if e[3] == 0 && e[2] >= mid {
          dsu.union(e[0] as usize, e[1] as usize);
        }
      }

      // Paid optional edges (require one upgrade, 2*s >= mid)
      let mut upgrades = 0i32;
      for e in &edges {
        if e[3] == 0 && e[2] < mid && e[2] * 2 >= mid {
          let pu = dsu.find(e[0] as usize);
          let pv = dsu.find(e[1] as usize);
          if pu != pv && upgrades < k {
            dsu.union(e[0] as usize, e[1] as usize);
            upgrades += 1;
          }
        }
      }

      dsu.components == 1
    };

    if !check(1) {
      return -1;
    }

    let mut lo = 1i32;
    let mut hi = max_s * 2;
    let mut ans = -1i32;

    while lo <= hi {
      let mid = lo + (hi - lo) / 2;
      if check(mid) {
        ans = mid;
        lo = mid + 1;
      } else {
        hi = mid - 1;
      }
    }

    ans
  }
}