Number of nodes in the sub tree with the same label

Hash Table Tree Depth-First Search Breadth-First Search Counting

55.4% acceptance

Feb 25, 2026

2370

813

You are given a tree (i.e. a connected, undirected graph that has no cycles) consisting of n nodes numbered from 0 to n - 1 and exactly n - 1 edges. The root of the tree is the node 0, and each node of the tree has a label which is a lower-case character given in the string labels (i.e. The node with the number i has the label labels[i]). The edges array is given on the form edges[i] = [ai, bi], which means there is an edge between nodes ai and bi in the tree. Return an array of size n where ans[i] is the number of nodes in the subtree of the ith node which have the same label as node i. A subtree of a tree T is the tree consisting of a node in T and all of its descendant nodes.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn count_sub_trees(n: i32, edges: Vec<Vec<i32>>, labels: String) -> Vec<i32> {
    let n = n as usize;
    let lb = labels.as_bytes();
    let mut adj: Vec<Vec<usize>> = vec![vec![]; n];
    for e in &edges {
      let (a, b) = (e[0] as usize, e[1] as usize);
      adj[a].push(b);
      adj[b].push(a);
    }
    let mut ans = vec![0i32; n];
    // Iterative DFS with post-order processing
    let mut counts: Vec<[i32; 26]> = vec![[0; 26]; n];
    let mut stack: Vec<(usize, usize, bool)> = vec![(0, usize::MAX, false)];
    while let Some((node, parent, visited)) = stack.pop() {
      if visited {
        let ch = (lb[node] - b'a') as usize;
        counts[node][ch] += 1;
        ans[node] = counts[node][ch];
        if parent != usize::MAX {
          for i in 0..26 {
            counts[parent][i] += counts[node][i];
          }
        }
      } else {
        stack.push((node, parent, true));
        for &nb in &adj[node] {
          if nb != parent {
            stack.push((nb, node, false));
          }
        }
      }
    }
    ans
  }
}