Count Subtrees With Max Distance Between Cities

hard trees bitmask floyd-warshall

Problem

There are n cities (1..n) that form a tree given by n-1 bidirectional edges. A subtree is a connected subset of cities. For each d from 1 to n-1, count the subtrees whose maximum distance between any two cities (number of edges on the path) equals exactly d. Return an array of size n-1 where the d-th element (1-indexed) is that count.

Inputn = 4, edges = [[1,2],[2,3],[2,4]]

Output[3,4,0]

{1,2}, {2,3}, {2,4} have diameter 1. {1,2,3}, {1,2,4}, {2,3,4}, {1,2,3,4} have diameter 2. None have diameter 3.

Inputn = 3, edges = [[1,2],[2,3]]

Output[2,1]

{1,2} and {2,3} give diameter 1; {1,2,3} is a path with diameter 2.

n (number of cities)

edges (u-v, comma-separated)

def count_subtrees(n, edges):
    INF = n + 1
    dist = [[INF] * n for _ in range(n)]
    for i in range(n):
        dist[i][i] = 0
    for u, v in edges:
        dist[u - 1][v - 1] = dist[v - 1][u - 1] = 1
    for k in range(n):
        for i in range(n):
            for j in range(n):
                if dist[i][k] + dist[k][j] < dist[i][j]:
                    dist[i][j] = dist[i][k] + dist[k][j]
    ans = [0] * (n - 1)
    for mask in range(1, 1 << n):
        nodes = [i for i in range(n) if mask & (1 << i)]
        if len(nodes) < 2:
            continue
        ec = sum(1 for u, v in edges
                 if mask & (1 << (u - 1)) and mask & (1 << (v - 1)))
        if ec != len(nodes) - 1:
            continue
        diam = max(dist[i][j] for i in nodes for j in nodes)
        ans[diam - 1] += 1
    return ans

function countSubtrees(n, edges) {
  const INF = n + 1;
  const dist = Array.from({ length: n }, () => Array(n).fill(INF));
  for (let i = 0; i < n; i++) dist[i][i] = 0;
  for (const [u, v] of edges) {
    dist[u - 1][v - 1] = dist[v - 1][u - 1] = 1;
  }
  for (let k = 0; k < n; k++)
    for (let i = 0; i < n; i++)
      for (let j = 0; j < n; j++)
        if (dist[i][k] + dist[k][j] < dist[i][j])
          dist[i][j] = dist[i][k] + dist[k][j];
  const ans = new Array(n - 1).fill(0);
  for (let mask = 1; mask < (1 << n); mask++) {
    const nodes = [];
    for (let i = 0; i < n; i++) if (mask & (1 << i)) nodes.push(i);
    if (nodes.length < 2) continue;
    let ec = 0;
    for (const [u, v] of edges)
      if ((mask & (1 << (u - 1))) && (mask & (1 << (v - 1)))) ec++;
    if (ec !== nodes.length - 1) continue;
    let diam = 0;
    for (const i of nodes) for (const j of nodes) diam = Math.max(diam, dist[i][j]);
    ans[diam - 1]++;
  }
  return ans;
}

int[] countSubtrees(int n, int[][] edges) {
    int INF = n + 1;
    int[][] dist = new int[n][n];
    for (int[] row : dist) java.util.Arrays.fill(row, INF);
    for (int i = 0; i < n; i++) dist[i][i] = 0;
    for (int[] e : edges)
        dist[e[0] - 1][e[1] - 1] = dist[e[1] - 1][e[0] - 1] = 1;
    for (int k = 0; k < n; k++)
        for (int i = 0; i < n; i++)
            for (int j = 0; j < n; j++)
                if (dist[i][k] + dist[k][j] < dist[i][j])
                    dist[i][j] = dist[i][k] + dist[k][j];
    int[] ans = new int[n - 1];
    for (int mask = 1; mask < (1 << n); mask++) {
        java.util.List<Integer> nodes = new java.util.ArrayList<>();
        for (int i = 0; i < n; i++) if ((mask & (1 << i)) != 0) nodes.add(i);
        if (nodes.size() < 2) continue;
        int ec = 0;
        for (int[] e : edges)
            if ((mask & (1 << (e[0] - 1))) != 0 && (mask & (1 << (e[1] - 1))) != 0) ec++;
        if (ec != nodes.size() - 1) continue;
        int diam = 0;
        for (int i : nodes) for (int j : nodes) diam = Math.max(diam, dist[i][j]);
        ans[diam - 1]++;
    }
    return ans;
}

vector<int> countSubtrees(int n, vector<vector<int>>& edges) {
    int INF = n + 1;
    vector<vector<int>> dist(n, vector<int>(n, INF));
    for (int i = 0; i < n; i++) dist[i][i] = 0;
    for (auto& e : edges)
        dist[e[0] - 1][e[1] - 1] = dist[e[1] - 1][e[0] - 1] = 1;
    for (int k = 0; k < n; k++)
        for (int i = 0; i < n; i++)
            for (int j = 0; j < n; j++)
                if (dist[i][k] + dist[k][j] < dist[i][j])
                    dist[i][j] = dist[i][k] + dist[k][j];
    vector<int> ans(n - 1, 0);
    for (int mask = 1; mask < (1 << n); mask++) {
        vector<int> nodes;
        for (int i = 0; i < n; i++) if (mask & (1 << i)) nodes.push_back(i);
        if ((int)nodes.size() < 2) continue;
        int ec = 0;
        for (auto& e : edges)
            if ((mask & (1 << (e[0] - 1))) && (mask & (1 << (e[1] - 1)))) ec++;
        if (ec != (int)nodes.size() - 1) continue;
        int diam = 0;
        for (int i : nodes) for (int j : nodes) diam = max(diam, dist[i][j]);
        ans[diam - 1]++;
    }
    return ans;
}

Explanation

Because n ≤ 15, we can afford to look at every subset of cities — there are at most 2¹⁵ of them. A subset is a valid subtree only if it is connected; for a connected subset we measure its diameter (the largest distance between any two members) and increment the tally for that diameter.

First we precompute all-pairs distances with Floyd–Warshall. We seed dist[i][i] = 0 and dist[u][v] = 1 for every edge, then relax dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j]) over all intermediate k. In a tree the unique path length between two cities is exactly this shortest distance.

We enumerate subsets as bitmasks 1 .. 2ⁿ-1. For each mask we collect its member nodes. Subsets of size 0 or 1 cannot have a positive diameter, so we skip them.

To test connectivity we count how many tree edges have both endpoints inside the mask. A subset of k nodes is a connected subtree exactly when it contains k-1 such edges — fewer means it is disconnected, and more is impossible inside a tree. This edge-count test is a clean, fast connectivity check that avoids a separate traversal.

For each connected subset we take the maximum dist[i][j] over its members as the diameter d, then do ans[d-1] += 1. After scanning all masks, ans holds the answer for every d from 1 to n-1.

Time: O(2ⁿ · n²) Space: O(n²)