Number of Possible Sets of Closing Branches

Hard

#2639Time: O(2^n * n * (E + n log n)) Where `E` is `roads.length`. We iterate through `2^n` subsets. For each subset of size `s` (where `s <= n`), we run Dijkstra `s` times. Each run of Dijkstra on the subgraph takes `O(E_sub + s log s)`, where `E_sub` is the number of edges in the subgraph. In the worst case, this gives a complexity of `O(2^n * n * (roads.length + n log n))`.Space: O(n + roads.length) This is for storing the graph's adjacency list. Inside the loop, Dijkstra's algorithm requires O(n) additional space for the distance map and priority queue.1 company

Patterns

Bit Manipulation Enumeration

Algorithms

Shortest Path

Data structures

Heap (Priority Queue)Graph

Companies

Atlassian

Prompt

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.

The distance between two branches is the minimum total traveled length needed to reach one branch from another.

You are given integers n, maxDistance, and a 0-indexed 2D array roads, where roads[i] = [u_i, v_i, w_i] represents the undirected road between branches u_i and v_i with length w_i.

Return the number of possible sets of closing branches, so that any branch has a distance of at most maxDistance from any other.

Note that, after closing a branch, the company will no longer have access to any roads connected to it.

Note that, multiple roads are allowed.

Example 1:

Input: n = 3, maxDistance = 5, roads = [[0,1,2],[1,2,10],[0,2,10]]
Output: 5
Explanation: The possible sets of closing branches are:
- The set [2], after closing, active branches are [0,1] and they are reachable to each other within distance 2.
- The set [0,1], after closing, the active branch is [2].
- The set [1,2], after closing, the active branch is [0].
- The set [0,2], after closing, the active branch is [1].
- The set [0,1,2], after closing, there are no active branches.
It can be proven, that there are only 5 possible sets of closing branches.

Example 2:

Input: n = 3, maxDistance = 5, roads = [[0,1,20],[0,1,10],[1,2,2],[0,2,2]]
Output: 7
Explanation: The possible sets of closing branches are:
- The set [], after closing, active branches are [0,1,2] and they are reachable to each other within distance 4.
- The set [0], after closing, active branches are [1,2] and they are reachable to each other within distance 2.
- The set [1], after closing, active branches are [0,2] and they are reachable to each other within distance 2.
- The set [0,1], after closing, the active branch is [2].
- The set [1,2], after closing, the active branch is [0].
- The set [0,2], after closing, the active branch is [1].
- The set [0,1,2], after closing, there are no active branches.
It can be proven, that there are only 7 possible sets of closing branches.

Example 3:

Input: n = 1, maxDistance = 10, roads = []
Output: 2
Explanation: The possible sets of closing branches are:
- The set [], after closing, the active branch is [0].
- The set [0], after closing, there are no active branches.
It can be proven, that there are only 2 possible sets of closing branches.

Constraints:

1 <= n <= 10
1 <= maxDistance <= 10⁵
0 <= roads.length <= 1000
roads[i].length == 3
0 <= u_i, v_i <= n - 1
u_i != v_i
1 <= w_i <= 1000
All branches are reachable from each other by traveling some roads.

Approaches

2 approaches with complexity analysis and trade-offs.

The problem asks for the number of valid sets of active branches. Since the number of branches n is very small (n <= 10), we can iterate through all 2^n possible subsets of branches. For each subset, we must verify if it's a valid set.

A set is valid if the shortest distance between any two branches in the set is at most maxDistance, considering only paths that use other branches from the same set. This approach checks the validity of each subset by running a single-source shortest path algorithm. For every node in the subset, we can run Dijkstra's algorithm to find its distances to all other nodes in the same subset and check if they meet the maxDistance criteria.

Algorithm

Initialize count = 0.
Iterate through all 2^n possible subsets of branches using a bitmask from 0 to 2^n - 1.
For each subset (represented by a mask):
- Create a list of activeNodes present in the current subset.
- If the subset is empty or contains a single node, it's trivially valid. Check this condition and continue.
- Assume the subset is valid by setting a flag is_valid = true.
- For each startNode in activeNodes:
  1. Run Dijkstra's algorithm starting from startNode. The search should be constrained to only traverse through other nodes that are also in the current activeNodes set.
  2. This will compute the shortest distances from startNode to all other reachable active nodes.
  3. Check if any computed distance to another endNode in activeNodes is greater than maxDistance.
  4. If such a distance is found, the subset is invalid. Set is_valid = false and break from the loops for the current subset.
- If is_valid remains true after checking all pairs, it means all nodes in the subset are within maxDistance of each other. Increment the count.
Return the total count.

Walkthrough

The core idea is to perform an exhaustive search over all 2^n subsets of branches. For each subset, we form a subgraph containing only the active branches and the roads connecting them. Then, we must verify the distance constraint. A simple way to do this is to compute all-pairs shortest paths within this subgraph. This can be achieved by running Dijkstra's algorithm from every single node in the active set.

For a given subset, we iterate through each of its nodes, treating it as a source. We run Dijkstra's to find the shortest paths to all other nodes in the subset. During the path search, we must ensure that we only traverse through nodes that are part of the current active subset. If at any point a shortest distance between two nodes in the set is found to be greater than maxDistance, we can immediately conclude that this subset is not valid and move on to the next one. If we check all pairs and all distances are within the limit, we count the subset as a valid one.

1import java.util.*;2 3class Solution {4    public int numberOfGoodSets(int n, int maxDistance, int[][] roads) {5        List<int[]>[] adj = new ArrayList[n];6        for (int i = 0; i < n; i++) {7            adj[i] = new ArrayList<>();8        }9        for (int[] road : roads) {10            adj[road[0]].add(new int[]{road[1], road[2]});11            adj[road[1]].add(new int[]{road[0], road[2]});12        }13 14        int goodSetsCount = 0;15        for (int mask = 0; mask < (1 << n); mask++) {16            if (isValid(mask, n, maxDistance, adj)) {17                goodSetsCount++;18            }19        }20        return goodSetsCount;21    }22 23    private boolean isValid(int mask, int n, int maxDistance, List<int[]>[] adj) {24        List<Integer> activeNodes = new ArrayList<>();25        for (int i = 0; i < n; i++) {26            if ((mask & (1 << i)) != 0) {27                activeNodes.add(i);28            }29        }30 31        if (activeNodes.size() <= 1) {32            return true;33        }34 35        for (int startNode : activeNodes) {36            Map<Integer, Integer> dists = dijkstra(startNode, mask, adj);37            for (int endNode : activeNodes) {38                if (!dists.containsKey(endNode) || dists.get(endNode) > maxDistance) {39                    return false;40                }41            }42        }43        return true;44    }45 46    private Map<Integer, Integer> dijkstra(int startNode, int mask, List<int[]>[] adj) {47        Map<Integer, Integer> dist = new HashMap<>();48        dist.put(startNode, 0);49 50        PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(a -> a[1]));51        pq.offer(new int[]{startNode, 0});52 53        while (!pq.isEmpty()) {54            int[] current = pq.poll();55            int u = current[0];56            int d = current[1];57 58            if (d > dist.getOrDefault(u, Integer.MAX_VALUE)) {59                continue;60            }61 62            for (int[] edge : adj[u]) {63                int v = edge[0];64                int weight = edge[1];65                if ((mask & (1 << v)) != 0) { // Only consider nodes in the subset66                    if (dist.getOrDefault(u, Integer.MAX_VALUE) + weight < dist.getOrDefault(v, Integer.MAX_VALUE)) {67                        dist.put(v, d + weight);68                        pq.offer(new int[]{v, d + weight});69                    }70                }71            }72        }73        return dist;74    }75}

Complexity

Time

O(2^n * n * (E + n log n)) Where `E` is `roads.length`. We iterate through `2^n` subsets. For each subset of size `s` (where `s <= n`), we run Dijkstra `s` times. Each run of Dijkstra on the subgraph takes `O(E_sub + s log s)`, where `E_sub` is the number of edges in the subgraph. In the worst case, this gives a complexity of `O(2^n * n * (roads.length + n log n))`.

Space

O(n + roads.length) This is for storing the graph's adjacency list. Inside the loop, Dijkstra's algorithm requires O(n) additional space for the distance map and priority queue.

Trade-offs

Pros

Conceptually straightforward, breaking the problem down into iterating through subsets and applying a standard graph algorithm (Dijkstra).
Correctly solves the problem within the given constraints.

Cons

Less efficient than the Floyd-Warshall approach for this problem's constraints due to the overhead of repeatedly initializing and running Dijkstra's algorithm.
The complexity of running Dijkstra from each node in the subset can be higher than a single run of Floyd-Warshall, especially given the small value of n.

This approach also iterates through all 2^n subsets of branches. However, instead of running a single-source shortest path algorithm multiple times for each subset, it uses the Floyd-Warshall algorithm to efficiently calculate all-pairs shortest paths (APSP) within the subgraph induced by the active branches. Floyd-Warshall is particularly well-suited for this problem because n is small and it computes all shortest paths simultaneously with a compact implementation.

Algorithm

First, create a base n x n adjacency matrix dist representing the direct distances between branches. Initialize dist[i][j] to a large value (infinity) if there's no direct road, 0 if i == j, and w for a road (i, j) with weight w. Handle multiple roads between the same two branches by taking the one with the minimum weight.
Initialize count = 0.
Iterate through all 2^n subsets using a bitmask from 0 to 2^n - 1.
For each subset (mask):
- Create a temporary copy of the base distance matrix for this subset's calculation.
- Apply the Floyd-Warshall algorithm to this temporary matrix. The key modification is that the intermediate node k, the source i, and the destination j must all be part of the current subset (i.e., their corresponding bits in the mask must be set).
- After the algorithm completes, check the validity of the subset. Iterate through all pairs of nodes (i, j) that are in the subset. If dist[i][j] > maxDistance for any pair, the subset is invalid.
- If the subset is valid after checking all its pairs, increment count.
Return count.

Walkthrough

We begin by pre-processing the roads into an n x n adjacency matrix, which will serve as a template for each subset's distance calculations. Then, we loop through every possible subset of branches, represented by a bitmask.

For each subset, we run a modified version of the Floyd-Warshall algorithm. The standard algorithm iterates through all possible intermediate nodes k to find shorter paths. In our modified version, we only consider nodes k, i, and j that are part of the current active subset. This ensures that the calculated shortest paths only use routes and intermediate branches that are currently active.

After running the algorithm for a subset, the resulting distance matrix contains the shortest paths between all pairs of active nodes within that subset's context. We then perform a final check: iterate through all pairs of active nodes (i, j) and verify if their shortest distance dist[i][j] is within maxDistance. If all pairs satisfy this condition, we've found a valid set, and we increment our counter.

1import java.util.Arrays;2 3class Solution {4    public int numberOfGoodSets(int n, int maxDistance, int[][] roads) {5        long[][] baseDist = new long[n][n];6        long infinity = (long) 1e11; // A large enough value for infinity7 8        for (int i = 0; i < n; i++) {9            Arrays.fill(baseDist[i], infinity);10            baseDist[i][i] = 0;11        }12 13        for (int[] road : roads) {14            int u = road[0];15            int v = road[1];16            int w = road[2];17            baseDist[u][v] = Math.min(baseDist[u][v], w);18            baseDist[v][u] = Math.min(baseDist[v][u], w);19        }20 21        int goodSetsCount = 0;22        for (int mask = 0; mask < (1 << n); mask++) {23            long[][] dist = new long[n][n];24            for (int i = 0; i < n; i++) {25                dist[i] = Arrays.copyOf(baseDist[i], n);26            }27 28            // Floyd-Warshall on the subgraph defined by the mask29            for (int k = 0; k < n; k++) {30                if ((mask & (1 << k)) != 0) { // k is in the subset31                    for (int i = 0; i < n; i++) {32                        if ((mask & (1 << i)) != 0) { // i is in the subset33                            for (int j = 0; j < n; j++) {34                                if ((mask & (1 << j)) != 0) { // j is in the subset35                                    dist[i][j] = Math.min(dist[i][j], dist[i][k] + dist[k][j]);36                                }37                            }38                        }39                    }40                }41            }42 43            boolean isValid = true;44            for (int i = 0; i < n; i++) {45                if ((mask & (1 << i)) != 0) {46                    for (int j = i; j < n; j++) {47                        if ((mask & (1 << j)) != 0) {48                            if (dist[i][j] > maxDistance) {49                                isValid = false;50                                break;51                            }52                        }53                    }54                }55                if (!isValid) break;56            }57 58            if (isValid) {59                goodSetsCount++;60            }61        }62        return goodSetsCount;63    }64}

Complexity

Time

O(2^n * n^3) The outer loop runs `2^n` times for each subset. Inside, the dominant operation is the modified Floyd-Warshall algorithm, which takes `O(n^3)` time, followed by a check that takes `O(n^2)`. The total complexity is determined by the product of these two factors.

Space

O(n^2) This space is used to store the base distance matrix and the temporary distance matrix for each subset's computation.

Trade-offs

Pros

More efficient than the Dijkstra-based approach for small n because it avoids the overhead of multiple separate pathfinding runs.
The implementation is clean and directly reflects the all-pairs shortest path nature of the validity check.
It is generally faster for small, dense graphs, which is effectively the case here.

Cons

The time complexity is exponential in n, making it suitable only for the small values of n given in the constraints.
Requires O(n^2) space, which is more than the adjacency list representation, but acceptable for n <= 10.

Solutions

Solution

1class Solution {2public3  int numberOfSets(int n, int maxDistance, int[][] roads) {4    int ans = 0;5    for (int mask = 0; mask < 1 << n; ++mask) {6      int[][] g = new int[n][n];7      for (var e : g) {8        Arrays.fill(e, 1 << 29);9      }10      for (var e : roads) {11        int u = e[0], v = e[1], w = e[2];12        if ((mask >> u & 1) == 1 && (mask >> v & 1) == 1) {13          g[u][v] = Math.min(g[u][v], w);14          g[v][u] = Math.min(g[v][u], w);15        }16      }17      for (int k = 0; k < n; ++k) {18        if ((mask >> k & 1) == 1) {19          g[k][k] = 0;20          for (int i = 0; i < n; ++i) {21            for (int j = 0; j < n; ++j) {22              g[i][j] = Math.min(g[i][j], g[i][k] + g[k][j]);23            }24          }25        }26      }27      int ok = 1;28      for (int i = 0; i < n && ok == 1; ++i) {29        for (int j = 0; j < n && ok == 1; ++j) {30          if ((mask >> i & 1) == 1 && (mask >> j & 1) == 1) {31            if (g[i][j] > maxDistance) {32              ok = 0;33            }34          }35        }36      }37      ans += ok;38    }39    return ans;40  }41}

1class Solution { public: int numberOfSets ( int n , int maxDistance , vector < vector < int >>& roads ) { int ans = 0 ; for ( int mask = 0 ; mask < 1 << n ; ++ mask ) { int g [ n ][ n ]; memset ( g , 0x3f , sizeof ( g )); for ( auto & e : roads ) { int u = e [ 0 ], v = e [ 1 ], w = e [ 2 ]; if (( mask >> u & 1 ) & ( mask >> v & 1 )) { g [ u ][ v ] = min ( g [ u ][ v ], w ); g [ v ][ u ] = min ( g [ v ][ u ], w ); } } for ( int k = 0 ; k < n ; ++ k ) { if ( mask >> k & 1 ) { g [ k ][ k ] = 0 ; for ( int i = 0 ; i < n ; ++ i ) { for ( int j = 0 ; j < n ; ++ j ) { g [ i ][ j ] = min ( g [ i ][ j ], g [ i ][ k ] + g [ k ][ j ]); } } } } int ok = 1 ; for ( int i = 0 ; i < n && ok == 1 ; ++ i ) { for ( int j = 0 ; j < n && ok == 1 ; ++ j ) { if (( mask >> i & 1 ) & ( mask >> j & 1 ) && g [ i ][ j ] > maxDistance ) { ok = 0 ; } } } ans += ok ; } return ans ; } };

1class Solution : def numberOfSets ( self , n : int , maxDistance : int , roads : List [ List [ int ]]) -> int : ans = 0 for mask in range ( 1 << n ): g = [[ inf ] * n for _ in range ( n )] for u , v , w in roads : if mask >> u & 1 and mask > v & 1 : g [ u ][ v ] = min ( g [ u ][ v ], w ) g [ v ][ u ] = min ( g [ v ][ u ], w ) for k in range ( n ): if mask >> k & 1 : g [ k ][ k ] = 0 for i in range ( n ): for j in range ( n ): # g[i][j] = min(g[i][j], g[i][k] + g[k][j]) if g [ i ][ k ] + g [ k ][ j ] < g [ i ][ j ]: g [ i ][ j ] = g [ i ][ k ] + g [ k ][ j ] if all ( g [ i ][ j ] <= maxDistance for i in range ( n ) for j in range ( n ) if mask >> i & 1 and mask >> j & 1 ): ans += 1 return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem