# Count Subtrees With Max Distance Between Cities
**Difficulty:** HARD
[External](https://leetcode.com/problems/count-subtrees-with-max-distance-between-cities)
Canonical: https://scaleengineer.com/dsa/problems/count-subtrees-with-max-distance-between-cities
**Patterns:** [Dynamic Programming](https://scaleengineer.com/dsa/patterns/dynamic-programming), [Bit Manipulation](https://scaleengineer.com/dsa/patterns/bit-manipulation), [Enumeration](https://scaleengineer.com/dsa/patterns/enumeration), [Bitmask](https://scaleengineer.com/dsa/patterns/bitmask)
**Data structures:** Tree
---
## Problem
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.

**Example 1:**

**![](https://assets.glich.co/dsa/count-subtrees-with-max-distance-between-cities/image0.png)**

**Input:** n = 4, edges = [[1,2],[2,3],[2,4]]
**Output:** [3,4,0]
**Explanation:**
The subtrees with subsets {1,2}, {2,3} and {2,4} have a max distance of 1.
The subtrees with subsets {1,2,3}, {1,2,4}, {2,3,4} and {1,2,3,4} have a max distance of 2.
No subtree has two nodes where the max distance between them is 3.

**Example 2:**

**Input:** n = 2, edges = [[1,2]]
**Output:** [1]

**Example 3:**

**Input:** n = 3, edges = [[1,2],[2,3]]
**Output:** [2,1]

**Constraints:**

* `2 <= n <= 15`
* `edges.length == n-1`
* `edges[i].length == 2`
* `1 <= ui, vi <= n`
* All pairs `(ui, vi)` are distinct.

# Approaches
## Brute-Force with Naive Diameter Calculation
This approach systematically checks every possible subset of cities. For each subset, it first determines if it constitutes a valid subtree (i.e., the cities in the subset form a connected component). If it is a valid subtree, it then calculates the maximum distance (diameter) between any two cities within that subtree. The diameter is found by computing the shortest path between all pairs of nodes in the subtree and taking the maximum.
**Time:** O(n^2 * 2^n). There are `2^n` subsets. For a subset of size `k`, checking connectivity takes `O(k)`. Calculating the diameter takes `k` BFS runs, each taking `O(k)`, for a total of `O(k^2)`. The sum over all subsets `sum_{k=1 to n} (C(n, k) * O(k^2))` results in `O(n^2 * 2^n)`. · **Space:** O(n). The adjacency list requires O(n) space. The queue and visited set for BFS also take up to O(n) space.
**Pros:** Conceptually simple and directly follows the problem definition.; Correct for the given constraints, although not the most efficient.
**Cons:** The naive diameter calculation is inefficient, leading to a high time complexity.; For larger `n` (even slightly above 15), this approach would be too slow.
### Explanation
The core idea is to iterate through all `2^n - 1` non-empty subsets of the `n` cities. A bitmask is a convenient way to represent these subsets, where the `i`-th bit being set means city `i+1` is in the subset.

For each subset:
1.  **Connectivity Check**: We must verify that the chosen subset of cities forms a connected graph. We can do this by picking an arbitrary city from the subset and performing a graph traversal (like BFS or DFS), making sure to only visit other cities within the same subset. If the number of visited cities equals the size of the subset, it's a valid subtree.
2.  **Diameter Calculation**: If the subset is a valid subtree, we find its diameter. A straightforward way to do this is to find the shortest path between every pair of nodes `(u, v)` in the subtree. This can be done by running a separate BFS starting from each node `u` in the subtree. The maximum distance found among all pairs is the diameter.
3.  **Counting**: Once the diameter `d` is found for a subtree, we increment the count for that distance in our result array.

The final result is an array where the `d`-th element stores the total count of subtrees with a maximum distance of `d`.

```java
class Solution {
    private java.util.List<Integer>[] adj;
    private int n;

    public int[] countSubtreesWithMaxDistance(int n, int[][] edges) {
        this.n = n;
        this.adj = new java.util.ArrayList[n];
        for (int i = 0; i < n; i++) {
            adj[i] = new java.util.ArrayList<>();
        }
        for (int[] edge : edges) {
            int u = edge[0] - 1;
            int v = edge[1] - 1;
            adj[u].add(v);
            adj[v].add(u);
        }

        int[] ans = new int[n - 1];
        for (int i = 1; i < (1 << n); i++) {
            java.util.List<Integer> subset = new java.util.ArrayList<>();
            for (int j = 0; j < n; j++) {
                if ((i & (1 << j)) != 0) {
                    subset.add(j);
                }
            }

            if (subset.size() < 2) {
                continue;
            }

            // 1. Check connectivity
            int startNode = subset.get(0);
            java.util.Queue<Integer> q = new java.util.LinkedList<>();
            q.offer(startNode);
            java.util.Set<Integer> visited = new java.util.HashSet<>();
            visited.add(startNode);

            while (!q.isEmpty()) {
                int u = q.poll();
                for (int v : adj[u]) {
                    if (subset.contains(v) && !visited.contains(v)) {
                        visited.add(v);
                        q.offer(v);
                    }
                }
            }

            if (visited.size() != subset.size()) {
                continue; // Not a valid subtree
            }

            // 2. Calculate diameter
            int maxDist = 0;
            for (int u : subset) {
                maxDist = Math.max(maxDist, bfsToFindMaxDist(u, subset));
            }

            if (maxDist > 0) {
                ans[maxDist - 1]++;
            }
        }
        return ans;
    }

    private int bfsToFindMaxDist(int startNode, java.util.List<Integer> subset) {
        java.util.Queue<int[]> q = new java.util.LinkedList<>(); // {node, dist}
        q.offer(new int[]{startNode, 0});
        java.util.Map<Integer, Integer> distances = new java.util.HashMap<>();
        distances.put(startNode, 0);
        int maxD = 0;

        while (!q.isEmpty()) {
            int[] curr = q.poll();
            int u = curr[0];
            int d = curr[1];
            maxD = Math.max(maxD, d);

            for (int v : adj[u]) {
                if (subset.contains(v) && !distances.containsKey(v)) {
                    distances.put(v, d + 1);
                    q.offer(new int[]{v, d + 1});
                }
            }
        }
        return maxD;
    }
}
```
### Algorithm
- Build an adjacency list representation of the tree from the input `edges`.
- Initialize an answer array `ans` of size `n-1`.
- Iterate through all `2^n - 1` non-empty subsets of cities using a bitmask from `1` to `(1 << n) - 1`.
- For each subset:
  - Check if the subset of nodes forms a connected component (a subtree). This is done by a single BFS/DFS traversal starting from an arbitrary node in the subset.
  - If it is a valid subtree and has at least two nodes:
    - Initialize `max_diameter = 0`.
    - For each node `u` in the subtree, run a BFS starting from `u` to find the maximum shortest-path distance to any other node `v` in the same subtree.
    - Update `max_diameter` with the maximum distance found across all starting nodes.
    - If `max_diameter > 0`, increment `ans[max_diameter - 1]`.
- Return the `ans` array.

## Brute-Force with Efficient Diameter Calculation
This approach improves upon the previous one by using a more efficient algorithm to calculate the diameter of a tree. While still iterating through all possible subsets of cities and checking for connectivity, it finds the diameter of a valid subtree in linear time with respect to the subtree's size. This is achieved by performing only two Breadth-First Searches (BFS).
**Time:** O(n * 2^n). For a subset of size `k`, connectivity check is `O(k)`. The efficient diameter calculation involves two BFS runs, each taking `O(k)`, for a total of `O(k)`. The sum over all subsets `sum_{k=1 to n} (C(n, k) * O(k))` results in `O(n * 2^n)`. · **Space:** O(n). The space required is for the adjacency list, BFS queue, and visited set, all of which are proportional to `n`.
**Pros:** Significantly more efficient than the naive approach.; Passes within the time limits for the given constraints (`n <= 15`).; Still relatively straightforward to implement.
**Cons:** The time complexity is still exponential, making it unsuitable for larger values of `n`.
### Explanation
The overall structure is the same: iterate through all `2^n - 1` subsets of cities.

The connectivity check for each subset remains the same as in the previous approach.

The key improvement is in the diameter calculation for a valid subtree. The diameter of any tree (or a connected acyclic graph) can be found using the following two-step process:
1.  Start a BFS from an arbitrary node `u` in the subtree to find the node `v` that is farthest from `u`.
2.  Start a second BFS from this farthest node `v`. The maximum distance found in this second BFS is the diameter of the tree. The node `w` at this maximum distance from `v` forms a pair `(v, w)` that is one of the endpoints of a longest path.

This two-BFS method reduces the time to find the diameter of a `k`-node subtree from `O(k^2)` to `O(k)`.

```java
class Solution {
    private java.util.List<Integer>[] adj;
    private int n;

    public int[] countSubtreesWithMaxDistance(int n, int[][] edges) {
        this.n = n;
        this.adj = new java.util.ArrayList[n];
        for (int i = 0; i < n; i++) {
            adj[i] = new java.util.ArrayList<>();
        }
        for (int[] edge : edges) {
            int u = edge[0] - 1;
            int v = edge[1] - 1;
            adj[u].add(v);
            adj[v].add(u);
        }

        int[] ans = new int[n - 1];
        for (int i = 1; i < (1 << n); i++) {
            int subsetSize = Integer.bitCount(i);
            if (subsetSize < 2) {
                continue;
            }

            int startNode = -1;
            for (int j = 0; j < n; j++) {
                if (((i >> j) & 1) == 1) {
                    startNode = j;
                    break;
                }
            }

            // 1. Check connectivity
            java.util.Queue<Integer> q = new java.util.LinkedList<>();
            q.offer(startNode);
            java.util.Set<Integer> visited = new java.util.HashSet<>();
            visited.add(startNode);

            while (!q.isEmpty()) {
                int u = q.poll();
                for (int v : adj[u]) {
                    if (((i >> v) & 1) == 1 && !visited.contains(v)) {
                        visited.add(v);
                        q.offer(v);
                    }
                }
            }

            if (visited.size() != subsetSize) {
                continue; // Not a valid subtree
            }

            // 2. Efficiently calculate diameter
            // First BFS to find one endpoint of a diameter
            int[] res1 = bfsForDiameter(startNode, i);
            int farthestNode = res1[0];
            
            // Second BFS from the endpoint to find the diameter
            int[] res2 = bfsForDiameter(farthestNode, i);
            int diameter = res2[1];

            if (diameter > 0) {
                ans[diameter - 1]++;
            }
        }
        return ans;
    }

    // Returns {farthest_node, max_distance}
    private int[] bfsForDiameter(int startNode, int mask) {
        java.util.Queue<int[]> q = new java.util.LinkedList<>(); // {node, dist}
        q.offer(new int[]{startNode, 0});
        java.util.Set<Integer> visited = new java.util.HashSet<>();
        visited.add(startNode);
        
        int farthestNode = startNode;
        int maxDist = 0;

        while (!q.isEmpty()) {
            int[] curr = q.poll();
            int u = curr[0];
            int d = curr[1];

            if (d > maxDist) {
                maxDist = d;
                farthestNode = u;
            }

            for (int v : adj[u]) {
                if (((mask >> v) & 1) == 1 && !visited.contains(v)) {
                    visited.add(v);
                    q.offer(new int[]{v, d + 1});
                }
            }
        }
        return new int[]{farthestNode, maxDist};
    }
}
```
### Algorithm
- Build an adjacency list for the tree.
- Initialize an answer array `ans`.
- Iterate through all `2^n - 1` non-empty subsets of cities.
- For each subset:
  - Verify connectivity using a single BFS/DFS traversal.
  - If it's a valid subtree with at least two nodes:
    - Pick an arbitrary node `u` from the subtree.
    - Run a BFS starting from `u` to find the farthest node `v`.
    - Run a second BFS starting from `v` to find its farthest node and the distance `d` to it. This distance `d` is the diameter.
    - If `d > 0`, increment `ans[d - 1]`.
- Return the `ans` array.

# Solutions
### Java

```java
class Solution {
private
  List<Integer>[] g;
private
  int msk;
private
  int nxt;
private
  int mx;
public
  int[] countSubgraphsForEachDiameter(int n, int[][] edges) {
    g = new List[n];
    Arrays.setAll(g, k->new ArrayList<>());
    for (int[] e : edges) {
      int u = e[0] - 1, v = e[1] - 1;
      g[u].add(v);
      g[v].add(u);
    }
    int[] ans = new int[n - 1];
    for (int mask = 1; mask < 1 << n; ++mask) {
      if ((mask & (mask - 1)) == 0) {
        continue;
      }
      msk = mask;
      mx = 0;
      int cur = 31 - Integer.numberOfLeadingZeros(msk);
      dfs(cur, 0);
      if (msk == 0) {
        msk = mask;
        mx = 0;
        dfs(nxt, 0);
        ++ans[mx - 1];
      }
    }
    return ans;
  }
private
  void dfs(int u, int d) {
    msk ^= 1 << u;
    if (mx < d) {
      mx = d;
      nxt = u;
    }
    for (int v : g[u]) {
      if ((msk >> v & 1) == 1) {
        dfs(v, d + 1);
      }
    }
  }
}

```

### CPP

```cpp
class Solution {
public:
  vector<int> countSubgraphsForEachDiameter(int n, vector<vector<int>> &edges) {
    vector<vector<int>> g(n);
    for (auto &e : edges) {
      int u = e[0] - 1, v = e[1] - 1;
      g[u].emplace_back(v);
      g[v].emplace_back(u);
    }
    vector<int> ans(n - 1);
    int nxt = 0, msk = 0, mx = 0;
    function<void(int, int)> dfs = [&](int u, int d) {
      msk ^= 1 << u;
      if (mx < d) {
        mx = d;
        nxt = u;
      }
      for (int &v : g[u]) {
        if (msk >> v & 1) {
          dfs(v, d + 1);
        }
      }
    };
    for (int mask = 1; mask < 1 << n; ++mask) {
      if ((mask & (mask - 1)) == 0) {
        continue;
      }
      msk = mask;
      mx = 0;
      int cur = 31 - __builtin_clz(msk);
      dfs(cur, 0);
      if (msk == 0) {
        msk = mask;
        mx = 0;
        dfs(nxt, 0);
        ++ans[mx - 1];
      }
    }
    return ans;
  }
};

```

### Python

```python
class Solution:
    def countSubgraphsForEachDiameter(self, n: int, edges: List[List[int]]) -> List[int]: def dfs(u: int, d: int = 0): nonlocal mx, nxt, msk if mx < d: mx, nxt = d, u msk ^= 1 << u for v in g[u]: if msk >> v & 1: dfs(v, d + 1) g = defaultdict(list) for u, v in edges: u, v = u - 1, v - 1 g[u]. append(v) g[v]. append(u) ans = [0] * (n - 1) nxt = mx = 0 for mask in range(1, 1 << n): if mask & (mask - 1) == 0: continue msk, mx = mask, 0 cur = msk . bit_length() - 1 dfs(cur) if msk == 0: msk, mx = mask, 0 dfs(nxt) ans[mx - 1] += 1 return ans

```