Count the number of good nodes

Tree Depth-First Search

55.3% acceptance

Feb 25, 2026

178

There is an undirected tree with n nodes labeled from 0 to n - 1, and rooted at node 0. You are given a 2D integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree. A node is good if all the subtrees rooted at its children have the same size. Return the number of good nodes in the given tree.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn count_good_nodes(edges: Vec<Vec<i32>>) -> i32 {
    let n = edges.len() + 1;
    let mut adj = vec![vec![]; n];
    for e in &edges {
      adj[e[0] as usize].push(e[1] as usize);
      adj[e[1] as usize].push(e[0] as usize);
    }
    let mut good = 0i32;
    // Iterative DFS
    let mut sub_size = vec![1usize; n];
    let mut parent = vec![n; n]; // n = no parent
    let mut order = Vec::with_capacity(n);
    let mut stack = vec![0usize];
    parent[0] = n;
    while let Some(v) = stack.pop() {
      order.push(v);
      for &u in &adj[v] {
        if parent[u] == n && u != 0 {
          parent[u] = v;
          stack.push(u);
        }
      }
    }
    // Process leaves first (reverse BFS order)
    for &v in order.iter().rev() {
      let children: Vec<usize> = adj[v].iter().copied().filter(|&u| parent[u] == v).collect();
      if children.is_empty() {
        // leaf: always good
        good += 1;
        continue;
      }
      let first_size = sub_size[children[0]];
      let is_good = children.iter().all(|&c| sub_size[c] == first_size);
      if is_good {
        good += 1;
      }
      let total: usize = children.iter().map(|&c| sub_size[c]).sum();
      sub_size[v] = total + 1;
    }
    good
  }
}