Number of possible sets of closing branches

Bit Manipulation Graph Theory Heap (Priority Queue) Enumeration Shortest Path

50.7% acceptance

Feb 25, 2026

201

There is a company with n branches across the country, some of which are connected by roads. Initially, all branches are reachable from each other by traveling some roads. The company has realized that they are spending an excessive amount of time traveling between their branches. As a result, they have decided to close down some of these branches (possibly none). However, they want to ensure that the remaining branches have a distance of at most maxDistance from each other. You are given integers n, maxDistance, and a 0-indexed 2D array roads, where roads[i] = [ui, vi, wi] represents the undirected road between branches ui and vi with length wi. Return the number of possible sets of closing branches, so that any branch has a distance of at most maxDistance from any other.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn number_of_sets(n: i32, max_distance: i32, roads: Vec<Vec<i32>>) -> i32 {
    let n = n as usize;
    let inf = i32::MAX / 2;
    let mut ans = 0;

    // Enumerate all 2^n subsets of open (active) branches
    for mask in 0u32..(1u32 << n) {
      // Run Floyd-Warshall on active nodes
      let mut dist = vec![vec![inf; n]; n];
      for i in 0..n {
        dist[i][i] = 0;
      }
      for road in &roads {
        let u = road[0] as usize;
        let v = road[1] as usize;
        let w = road[2];
        if (mask >> u) & 1 == 1 && (mask >> v) & 1 == 1 {
          dist[u][v] = dist[u][v].min(w);
          dist[v][u] = dist[v][u].min(w);
        }
      }
      for k in 0..n {
        if (mask >> k) & 1 == 0 { continue; }
        for i in 0..n {
          if (mask >> i) & 1 == 0 { continue; }
          for j in 0..n {
            if (mask >> j) & 1 == 0 { continue; }
            if dist[i][k] < inf && dist[k][j] < inf {
              dist[i][j] = dist[i][j].min(dist[i][k] + dist[k][j]);
            }
          }
        }
      }
      // Check: all active pairs have distance <= max_distance
      let mut valid = true;
      'outer: for i in 0..n {
        if (mask >> i) & 1 == 0 { continue; }
        for j in 0..n {
          if (mask >> j) & 1 == 0 { continue; }
          if dist[i][j] > max_distance {
            valid = false;
            break 'outer;
          }
        }
      }
      if valid { ans += 1; }
    }
    ans
  }
}