Problem: Number of Possible Sets of Closing Branches

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] = [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.

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.


Python solution:
import math
import itertools
import bisect
import collections
import string
import heapq
import functools
import sortedcontainers
from typing import List, Dict, Tuple, Iterator

class Solution:
  def numberOfSets(self, n: int, maxDistance: int, roads: List[List[int]]) -> int:
    return sum(self._floydWarshall(n, maxDistance, roads, mask) <= maxDistance for mask in range(1 << n))

  def _floydWarshall(self, n: int, maxDistanceThreshold: int, roads: List[List[int]], mask: int) -> List[List[int]]:
    maxDistance = 0
    dist = [[maxDistanceThreshold + 1] * n for _ in range(n)]

    for i in range(n):
      if mask >> i & 1:
        dist[i][i] = 0

    for u, v, w in roads:
      if mask >> u & 1 and mask >> v & 1:
        dist[u][v] = min(dist[u][v], w)
        dist[v][u] = min(dist[v][u], w)

    for k in range(n):
      if mask >> k & 1:
        for i in range(n):
          if mask >> i & 1:
            for j in range(n):
              if mask >> j & 1:
                dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])

    for i in range(n):
      if mask >> i & 1:
        for j in range(i + 1, n):
          if mask >> j & 1:
            maxDistance = max(maxDistance, dist[i][j])

    return maxDistance


Java solution:
class Solution {
  public int numberOfSets(int n, int maxDistance, int[][] roads) {
    final int maxMask = 1 << n;
    int ans = 0;

    for (int mask = 0; mask < maxMask; ++mask)
      if (floydWarshall(n, maxDistance, roads, mask) <= maxDistance)
        ++ans;

    return ans;
  }

  private int floydWarshall(int n, int maxDistanceThreshold, int[][] roads, int mask) {
    int maxDistance = 0;
    int[][] dist = new int[n][n];
    Arrays.stream(dist).forEach(A -> Arrays.fill(A, maxDistanceThreshold + 1));

    for (int i = 0; i < n; ++i)
      if ((mask >> i & 1) == 1)
        dist[i][i] = 0;

    for (int[] road : roads) {
      final int u = road[0];
      final int v = road[1];
      final int w = road[2];
      if ((mask >> u & 1) == 1 && (mask >> v & 1) == 1) {
        dist[u][v] = Math.min(dist[u][v], w);
        dist[v][u] = Math.min(dist[v][u], w);
      }
    }

    for (int k = 0; k < n; ++k)
      if ((mask >> k & 1) == 1)
        for (int i = 0; i < n; ++i)
          if ((mask >> i & 1) == 1)
            for (int j = 0; j < n; ++j)
              if ((mask >> j & 1) == 1)
                dist[i][j] = Math.min(dist[i][j], dist[i][k] + dist[k][j]);

    for (int i = 0; i < n; ++i)
      if ((mask >> i & 1) == 1)
        for (int j = i + 1; j < n; ++j)
          if ((mask >> j & 1) == 1)
            maxDistance = Math.max(maxDistance, dist[i][j]);

    return maxDistance;
  }
}


C++ solution:
class Solution {
 public:
  int numberOfSets(int n, int maxDistance, vector<vector<int>>& roads) {
    const int maxMask = 1 << n;
    int ans = 0;

    for (int mask = 0; mask < maxMask; ++mask)
      if (floydWarshall(n, maxDistance, roads, mask) <= maxDistance)
        ++ans;

    return ans;
  }

 private:
  // Returns the maximum distance between any two branches, where the mask
  // represents the selected branches.
  int floydWarshall(int n, int maxDistanceThreshold, vector<vector<int>>& roads,
                    int mask) {
    int maxDistance = 0;
    vector<vector<int>> dist(n, vector<int>(n, maxDistanceThreshold + 1));

    for (int i = 0; i < n; ++i)
      if (mask >> i & 1)
        dist[i][i] = 0;

    for (const vector<int>& road : roads) {
      const int u = road[0];
      const int v = road[1];
      const int w = road[2];
      if (mask >> u & 1 && mask >> v & 1) {
        dist[u][v] = min(dist[u][v], w);
        dist[v][u] = min(dist[v][u], w);
      }
    }

    for (int k = 0; k < n; ++k)
      if (mask >> k & 1)
        for (int i = 0; i < n; ++i)
          if (mask >> i & 1)
            for (int j = 0; j < n; ++j)
              if (mask >> j & 1)
                dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j]);

    for (int i = 0; i < n; ++i)
      if (mask >> i & 1)
        for (int j = i + 1; j < n; ++j)
          if (mask >> j & 1)
            maxDistance = max(maxDistance, dist[i][j]);

    return maxDistance;
  }
};