Problem: Count Subtrees With Max Distance Between Cities

There are `n` cities numbered from `1` to `n`. You are given an array `edges`
of size `n-1`, where `edges[i] = [ui, vi]` represents a bidirectional edge
between cities `ui` and `vi`. There exists a unique path between each pair of
cities. In other words, the cities form a tree.

A subtree is a subset of cities where every city is reachable from every other
city in the subset, where the path between each pair passes through only the
cities from the subset. Two subtrees are different if there is a city in one
subtree that is not present in the other.

For each `d` from `1` to `n-1`, find the number of subtrees in which the
maximum distance between any two cities in the subtree is equal to `d`.

Return an array of size `n-1` where the `dth` element (1-indexed) is the
number of subtrees in which the maximum distance between any two cities is
equal to `d`.

Notice that the distance between the two cities is the number of edges in the
path between them.


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 countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]:
    maxMask = 1 << n
    dist = self._floydWarshall(n, edges)
    ans = [0] * (n - 1)

    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:
        continue
      cityCount += 1
      for v in range(u + 1, n):
        if (mask >> v) & 1 == 0:
          continue
        if dist[u][v] == 1:
          edgeCount += 1
        maxDist = max(maxDist, dist[u][v])

    if edgeCount == cityCount - 1:
      return maxDist
    else:
      return 0


Java solution:
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;
  }
}


C++ solution:
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;
  }
};