You are given an integer n representing n cities numbered from 1 to n, and an array edges where each edges[i] = [u, v] indicates that there is a bidirectional road connecting city u and city v. The cities and roads form a connected, undirected tree.
For every integer d in the range 1 to n-1, you need to count how many connected subtrees of the original tree have a maximum distance between any two cities equal to d. The maximum distance between any two cities in a subtree is the length (number of edges) of the longest path in that subtree.
Return an array ans of length n-1 where ans[d-1] is the number of connected subtrees with max distance d.
At first glance, the problem seems to require examining all possible subsets of cities and checking which of them form connected subtrees. For each such subtree, we must compute the diameter (the largest distance between any two nodes).
The brute-force approach would be to generate all possible subsets of the n cities, check if each subset forms a connected tree, and then compute the diameter for each. However, since there are 2^n possible subsets, this is computationally infeasible for n up to 15.
We need an optimization. Since n is small (≤ 15), we can use bitmasking to represent subsets. For each subset, we can use BFS or DFS to check connectivity and calculate the diameter efficiently. The key insight is that the total number of subsets is manageable, and we can prune subsets that are not connected early.
n. If the ith bit is set, city i is included in the subset.
1 to 2^n - 1, consider only those with at least two cities (i.e., at least two bits set).
d, increment ans[d-1].
ans array.
We use adjacency lists for efficient graph traversal, and bitmasking to enumerate all possible city subsets. BFS is used both for connectivity checks and for computing the diameter within each subset.
Consider n = 4 and edges = [[1,2],[2,3],[2,4]]. The tree is:
1
|
2
/ \
3 4
ans = [3, 4, 0] (since n-1=3).
O(2^n \cdot n^2) — For each subset (of which there are 2^n), we check connectivity and compute the diameter (each can be done in O(n^2)).
n ≤ 15, this is acceptable (~32,000 subsets).
O(n^2) for the adjacency list and O(2^n) for bitmask enumeration.
The time complexity is exponential in n, but the constraints allow it.
The problem is solved by representing all possible city subsets as bitmasks, checking which ones are connected subtrees, and computing the diameter for each using BFS. The approach leverages the small value of n to perform an exhaustive search efficiently. The key insight is that for small trees, bitmask enumeration and BFS are practical and effective.
class Solution:
def countSubgraphsForEachDiameter(self, n, edges):
from collections import deque, defaultdict
# Build adjacency list
graph = [[] for _ in range(n)]
for u, v in edges:
graph[u-1].append(v-1)
graph[v-1].append(u-1)
ans = [0] * (n-1)
# For all subsets of nodes (bitmask representation)
for mask in range(1, 1 << n):
# Only consider subsets with at least 2 nodes
if bin(mask).count('1') < 2:
continue
# Find first node in subset
nodes = [i for i in range(n) if (mask & (1 << i))]
start = nodes[0]
# Check connectivity via BFS
visited = set()
queue = deque([start])
visited.add(start)
while queue:
curr = queue.popleft()
for nei in graph[curr]:
if (mask & (1 << nei)) and nei not in visited:
visited.add(nei)
queue.append(nei)
if len(visited) != len(nodes):
continue # Not connected
# Compute diameter in the subset
def bfs(u):
seen = set([u])
queue = deque([(u, 0)])
farthest = (u, 0)
while queue:
node, dist = queue.popleft()
if dist > farthest[1]:
farthest = (node, dist)
for v in graph[node]:
if (mask & (1 << v)) and v not in seen:
seen.add(v)
queue.append((v, dist+1))
return farthest
far_node, _ = bfs(start)
_, diameter = bfs(far_node)
if diameter > 0:
ans[diameter-1] += 1
return ans
class Solution {
public:
vector<int> countSubgraphsForEachDiameter(int n, vector<vector<int>>& edges) {
vector<vector<int>> graph(n);
for (auto& e : edges) {
int u = e[0] - 1, v = e[1] - 1;
graph[u].push_back(v);
graph[v].push_back(u);
}
vector<int> ans(n-1, 0);
int total = 1 << n;
for (int mask = 1; mask < total; ++mask) {
if (__builtin_popcount(mask) < 2) continue;
vector<int> nodes;
for (int i = 0; i < n; ++i)
if (mask & (1 << i)) nodes.push_back(i);
// Check connectivity via BFS
queue<int> q;
vector<bool> visited(n, false);
q.push(nodes[0]);
visited[nodes[0]] = true;
int count = 1;
while (!q.empty()) {
int curr = q.front(); q.pop();
for (int nei : graph[curr]) {
if ((mask & (1 << nei)) && !visited[nei]) {
visited[nei] = true;
q.push(nei);
++count;
}
}
}
if (count != nodes.size()) continue;
// Compute diameter in the subset
auto bfs = [&](int u) {
queue<pair<int,int>> q;
vector<bool> seen(n, false);
q.push({u,0});
seen[u] = true;
pair<int,int> farthest = {u,0};
while (!q.empty()) {
auto [node, dist] = q.front(); q.pop();
if (dist > farthest.second) farthest = {node, dist};
for (int v : graph[node]) {
if ((mask & (1 << v)) && !seen[v]) {
seen[v] = true;
q.push({v, dist+1});
}
}
}
return farthest;
};
auto far = bfs(nodes[0]);
auto far2 = bfs(far.first);
int diameter = far2.second;
if (diameter > 0) ans[diameter-1]++;
}
return ans;
}
};
class Solution {
public int[] countSubgraphsForEachDiameter(int n, int[][] edges) {
List<List<Integer>> graph = new ArrayList<>();
for (int i = 0; i < n; i++) graph.add(new ArrayList<>());
for (int[] e : edges) {
graph.get(e[0]-1).add(e[1]-1);
graph.get(e[1]-1).add(e[0]-1);
}
int[] ans = new int[n-1];
int total = 1 << n;
for (int mask = 1; mask < total; mask++) {
if (Integer.bitCount(mask) < 2) continue;
List<Integer> nodes = new ArrayList<>();
for (int i = 0; i < n; i++)
if ((mask & (1 << i)) != 0) nodes.add(i);
// Check connectivity via BFS
Queue<Integer> q = new LinkedList<>();
boolean[] visited = new boolean[n];
q.offer(nodes.get(0));
visited[nodes.get(0)] = true;
int count = 1;
while (!q.isEmpty()) {
int curr = q.poll();
for (int nei : graph.get(curr)) {
if (((mask & (1 << nei)) != 0) && !visited[nei]) {
visited[nei] = true;
q.offer(nei);
count++;
}
}
}
if (count != nodes.size()) continue;
// Compute diameter in the subset
class Pair {
int node, dist;
Pair(int n, int d) { node = n; dist = d; }
}
java.util.function.Function<Integer, Pair> bfs = (u) -> {
Queue<Pair> q2 = new LinkedList<>();
boolean[] seen = new boolean[n];
q2.offer(new Pair(u, 0));
seen[u] = true;
Pair farthest = new Pair(u, 0);
while (!q2.isEmpty()) {
Pair p = q2.poll();
if (p.dist > farthest.dist) farthest = p;
for (int v : graph.get(p.node)) {
if (((mask & (1 << v)) != 0) && !seen[v]) {
seen[v] = true;
q2.offer(new Pair(v, p.dist+1));
}
}
}
return farthest;
};
Pair far = bfs.apply(nodes.get(0));
Pair far2 = bfs.apply(far.node);
int diameter = far2.dist;
if (diameter > 0) ans[diameter-1]++;
}
return ans;
}
}
var countSubgraphsForEachDiameter = function(n, edges) {
const graph = Array.from({length: n}, () => []);
for (const [u, v] of edges) {
graph[u-1].push(v-1);
graph[v-1].push(u-1);
}
const ans = Array(n-1).fill(0);
const total = 1 << n;
for (let mask = 1; mask < total; ++mask) {
if (countBits(mask) < 2) continue;
const nodes = [];
for (let i = 0; i < n; ++i)
if (mask & (1 << i)) nodes.push(i);
// Check connectivity via BFS
const visited = new Array(n).fill(false);
const queue = [nodes[0]];
visited[nodes[0]] = true;
let count = 1;
while (queue.length) {
const curr = queue.shift();
for (const nei of graph[curr]) {
if ((mask & (1 << nei)) && !visited[nei]) {
visited[nei] = true;
queue.push(nei);
count++;
}
}
}
if (count !== nodes.length) continue;
// Compute diameter in the subset
function bfs(u) {
const seen = new Array(n).fill(false);
const queue = [[u, 0]];
seen[u] = true;
let farthest = [u, 0];
while (queue.length) {
const [node, dist] = queue.shift();
if (dist > farthest[1]) farthest = [node, dist];
for (const v of graph[node]) {
if ((mask & (1 << v)) && !seen[v]) {
seen[v] = true;
queue.push([v, dist+1]);
}
}
}
return farthest;
}
const far = bfs(nodes[0]);
const far2 = bfs(far[0]);
const diameter = far2[1];
if (diameter > 0) ans[diameter-1]++;
}
return ans;
function countBits(x) {
let cnt = 0;
while (x) {
cnt += x & 1;
x >>= 1;
}
return cnt;
}
};