Count subtrees with max distance between cities

Dynamic Programming Bit Manipulation Tree Enumeration Bitmask

67.4% acceptance

Feb 25, 2026

571

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.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn count_subgraphs_for_each_diameter(n: i32, edges: Vec<Vec<i32>>) -> Vec<i32> {
    let n = n as usize;
    // Build adjacency list (0-indexed)
    let mut adj = vec![vec![]; n];
    for e in &edges {
      let u = (e[0] - 1) as usize;
      let v = (e[1] - 1) as usize;
      adj[u].push(v);
      adj[v].push(u);
    }
    // BFS from each node to get all-pairs distances
    let mut dist = vec![vec![0usize; n]; n];
    for start in 0..n {
      let mut visited = vec![false; n];
      let mut queue = std::collections::VecDeque::new();
      queue.push_back(start);
      visited[start] = true;
      dist[start][start] = 0;
      while let Some(node) = queue.pop_front() {
        for &nb in &adj[node] {
          if !visited[nb] {
            visited[nb] = true;
            dist[start][nb] = dist[start][node] + 1;
            queue.push_back(nb);
          }
        }
      }
    }

    let mut ans = vec![0i32; n - 1];
    // Enumerate all non-empty subsets with >= 2 nodes
    for mask in 1u32..(1u32 << n) {
      if mask.count_ones() < 2 { continue; }
      let nodes: Vec<usize> = (0..n).filter(|&i| mask & (1 << i) != 0).collect();
      // Check connectivity: count edges within mask
      let mut edge_count = 0;
      for &u in &nodes {
        for &v in &adj[u] {
          if mask & (1 << v) != 0 { edge_count += 1; }
        }
      }
      edge_count /= 2;
      if edge_count != nodes.len() - 1 { continue; } // not connected subtree
      // Find max distance
      let mut max_d = 0;
      for i in 0..nodes.len() {
        for j in (i+1)..nodes.len() {
          max_d = max_d.max(dist[nodes[i]][nodes[j]]);
        }
      }
      if max_d > 0 { ans[max_d - 1] += 1; }
    }
    ans
  }
}