Find subtree sizes after changes

Array Hash Table String Tree Depth-First Search

54.0% acceptance

Feb 23, 2026

105

You are given a tree rooted at node 0 that consists of n nodes numbered from 0 to n - 1. The tree is represented by an array parent of size n, where parent[i] is the parent of node i. Since node 0 is the root, parent[0] == -1. You are also given a string s of length n, where s[i] is the character assigned to node i. We make the following changes on the tree one time simultaneously for all nodes x from 1 to n - 1: Find the closest node y to node x such that y is an ancestor of x, and s[x] == s[y]. If node y does not exist, do nothing. Otherwise, remove the edge between x and its current parent and make node y the new parent of x by adding an edge between them. Return an array answer of size n where answer[i] is the size of the subtree rooted at node i in the final tree.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn find_subtree_sizes(parent: Vec<i32>, s: String) -> Vec<i32> {
    let n = parent.len();
    let chars: Vec<u8> = s.bytes().collect();
    let mut children: Vec<Vec<usize>> = vec![vec![]; n];
    for i in 1..n {
      children[parent[i] as usize].push(i);
    }
    
    // DFS: track last seen ancestor with same character (stack per char)
    // When visiting node x: check if there's an ancestor with char s[x]
    // If so, change parent of x in the new tree to that ancestor
    
    let mut new_parent = parent.clone();
    let mut char_stack: Vec<Vec<usize>> = vec![vec![]; 26]; // stack per character
    
    // Iterative DFS with stack to simulate recursion and maintain ancestor char stacks
    let mut stack: Vec<(usize, bool)> = vec![(0, false)];
    
    while let Some((node, leaving)) = stack.pop() {
      if leaving {
        char_stack[(chars[node] - b'a') as usize].pop();
      } else {
        // Mark to cleanup when we leave
        stack.push((node, true));
        let ci = (chars[node] - b'a') as usize;
        if let Some(&ancestor) = char_stack[ci].last() {
          new_parent[node] = ancestor as i32;
        }
        char_stack[ci].push(node);
        // Push children (in reverse order to process in order)
        for &child in children[node].iter().rev() {
          stack.push((child, false));
        }
      }
    }
    
    // Compute subtree sizes using new_parent
    let mut size = vec![1i32; n];
    // Process nodes in reverse topological order (since root is 0, process leaves first)
    // We can traverse from n-1 downto 0 only if nodes are topologically ordered by number
    // which is not guaranteed. Instead, do a proper topological sort.
    // Use the new_parent array: root nodes contribute to parent.
    
    // Actually, just do DFS again on new tree
    let mut new_children: Vec<Vec<usize>> = vec![vec![]; n];
    for i in 1..n {
      new_children[new_parent[i] as usize].push(i);
    }
    
    // Iterative post-order DFS
    let mut order = Vec::with_capacity(n);
    let mut stk: Vec<usize> = vec![0];
    while let Some(node) = stk.pop() {
      order.push(node);
      for &c in &new_children[node] {
        stk.push(c);
      }
    }
    
    for &node in order.iter().rev() {
      if new_parent[node] != -1 {
        let p = new_parent[node] as usize;
        size[p] += size[node];
      }
    }
    
    size
  }
}