Skip to content

1617. Count Subtrees With Max Distance Between Cities 👍

  • Time: $O(2^n \cdot n^2)$
  • Space: $O(n^2)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
class Solution {
 public:
  vector<int> countSubgraphsForEachDiameter(int n, vector<vector<int>>& edges) {
    const int maxMask = 1 << n;
    const vector<vector<int>> dist = floydWarshall(n, edges);
    vector<int> ans(n - 1);

    // mask := the subset of the cities
    for (int mask = 0; mask < maxMask; ++mask) {
      const int maxDist = getMaxDist(mask, dist, n);
      if (maxDist > 0)
        ++ans[maxDist - 1];
    }

    return ans;
  }

 private:
  vector<vector<int>> floydWarshall(int n, const vector<vector<int>>& edges) {
    vector<vector<int>> dist(n, vector<int>(n, n));

    for (int i = 0; i < n; ++i)
      dist[i][i] = 0;

    for (const vector<int>& edge : edges) {
      const int u = edge[0] - 1;
      const int v = edge[1] - 1;
      dist[u][v] = 1;
      dist[v][u] = 1;
    }

    for (int k = 0; k < n; ++k)
      for (int i = 0; i < n; ++i)
        for (int j = 0; j < n; ++j)
          dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j]);

    return dist;
  }

  int getMaxDist(int mask, const vector<vector<int>>& dist, int n) {
    int maxDist = 0;
    int edgeCount = 0;
    int cityCount = 0;
    for (int u = 0; u < n; ++u) {
      if ((mask >> u & 1) == 0)  // u is not in the subset.
        continue;
      ++cityCount;
      for (int v = u + 1; v < n; ++v) {
        if ((mask >> v & 1) == 0)  // v is not in the subset.
          continue;
        if (dist[u][v] == 1)  // u and v are connected.
          ++edgeCount;
        maxDist = max(maxDist, dist[u][v]);
      }
    }
    return edgeCount == cityCount - 1 ? maxDist : 0;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
class Solution {
  public int[] countSubgraphsForEachDiameter(int n, int[][] edges) {
    final int maxMask = 1 << n;
    final int[][] dist = floydWarshall(n, edges);
    int[] ans = new int[n - 1];

    // mask := the subset of the cities
    for (int mask = 0; mask < maxMask; ++mask) {
      int maxDist = getMaxDist(mask, dist, n);
      if (maxDist > 0)
        ++ans[maxDist - 1];
    }

    return ans;
  }

  private int[][] floydWarshall(int n, int[][] edges) {
    int[][] dist = new int[n][n];
    for (int i = 0; i < n; ++i)
      Arrays.fill(dist[i], n);

    for (int i = 0; i < n; ++i)
      dist[i][i] = 0;

    for (int[] edge : edges) {
      final int u = edge[0] - 1;
      final int v = edge[1] - 1;
      dist[u][v] = 1;
      dist[v][u] = 1;
    }

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

    return dist;
  }

  private int getMaxDist(int mask, int[][] dist, int n) {
    int maxDist = 0;
    int edgeCount = 0;
    int cityCount = 0;
    for (int u = 0; u < n; ++u) {
      if ((mask >> u & 1) == 0) // u is not in the subset.
        continue;
      ++cityCount;
      for (int v = u + 1; v < n; ++v) {
        if ((mask >> v & 1) == 0) // v is not in the subset.
          continue;
        if (dist[u][v] == 1) // u and v are connected.
          ++edgeCount;
        maxDist = Math.max(maxDist, dist[u][v]);
      }
    }
    return edgeCount == cityCount - 1 ? maxDist : 0;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
class Solution:
  def countSubgraphsForEachDiameter(
      self,
      n: int,
      edges: list[list[int]],
  ) -> list[int]:
    maxMask = 1 << n
    dist = self._floydWarshall(n, edges)
    ans = [0] * (n - 1)

    # mask := the subset of the cities
    for mask in range(maxMask):
      maxDist = self._getMaxDist(mask, dist, n)
      if maxDist > 0:
        ans[maxDist - 1] += 1

    return ans

  def _floydWarshall(self, n: int, edges: list[list[int]]) -> list[list[int]]:
    dist = [[n] * n for _ in range(n)]

    for i in range(n):
      dist[i][i] = 0

    for u, v in edges:
      dist[u - 1][v - 1] = 1
      dist[v - 1][u - 1] = 1

    for k in range(n):
      for i in range(n):
        for j in range(n):
          dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])

    return dist

  def _getMaxDist(self, mask: int, dist: list[list[int]], n: int) -> int:
    maxDist = 0
    edgeCount = 0
    cityCount = 0
    for u in range(n):
      if (mask >> u) & 1 == 0:  # u is not in the subset.
        continue
      cityCount += 1
      for v in range(u + 1, n):
        if (mask >> v) & 1 == 0:  # v is not in the subset.
          continue
        if dist[u][v] == 1:  # u and v are connected.
          edgeCount += 1
        maxDist = max(maxDist, dist[u][v])

    return maxDist if edgeCount == cityCount - 1 else 0