Find Subtree Sizes After Changes

medium tree dfs hash map

Problem

A rooted tree on n nodes is given by parent (with parent[0] = -1), and each node i holds the letter s[i]. In one pass, for every non-root node x, if there is an ancestor with the same letter, detach x and reattach it under the closest such ancestor (all changes are decided from the original tree). Return an array where answer[i] is the size of the subtree rooted at i in the final tree.

Inputparent = [-1,0,0,1,1,1], s = "abaaba"

Output[6,2,1,1,1,1]

Nodes 3 and 5 (letter 'a') reattach to their closest 'a'-ancestor, the root 0, leaving node 1 with a single child.

parent array

def find_subtree_sizes(parent, s):
    n = len(parent)
    children = [[] for _ in range(n)]
    for i in range(1, n):
        children[parent[i]].append(i)
    new_parent = parent[:]
    stacks = {}
    def reassign(node):
        c = s[node]
        if stacks.get(c):
            new_parent[node] = stacks[c][-1]
        stacks.setdefault(c, []).append(node)
        for ch in children[node]:
            reassign(ch)
        stacks[c].pop()
    reassign(0)
    new_children = [[] for _ in range(n)]
    for i in range(1, n):
        new_children[new_parent[i]].append(i)
    size = [0] * n
    def measure(node):
        total = 1
        for ch in new_children[node]:
            total += measure(ch)
        size[node] = total
        return total
    measure(0)
    return size

function findSubtreeSizes(parent, s) {
  const n = parent.length;
  const children = Array.from({length: n}, () => []);
  for (let i = 1; i < n; i++) children[parent[i]].push(i);
  const newParent = parent.slice();
  const stacks = {};
  function reassign(node) {
    const c = s[node];
    if (stacks[c] && stacks[c].length) newParent[node] = stacks[c][stacks[c].length - 1];
    (stacks[c] = stacks[c] || []).push(node);
    for (const ch of children[node]) reassign(ch);
    stacks[c].pop();
  }
  reassign(0);
  const newChildren = Array.from({length: n}, () => []);
  for (let i = 1; i < n; i++) newChildren[newParent[i]].push(i);
  const size = new Array(n).fill(0);
  function measure(node) {
    let total = 1;
    for (const ch of newChildren[node]) total += measure(ch);
    size[node] = total;
    return total;
  }
  measure(0);
  return size;
}

import java.util.*;
class Solution {
    int[] newParent, size; List<Integer>[] children, newChildren;
    Map<Character, Deque<Integer>> stacks = new HashMap<>();
    String s;
    public int[] findSubtreeSizes(int[] parent, String s) {
        int n = parent.length; this.s = s;
        children = new List[n]; newChildren = new List[n];
        for (int i = 0; i < n; i++) { children[i] = new ArrayList<>(); newChildren[i] = new ArrayList<>(); }
        for (int i = 1; i < n; i++) children[parent[i]].add(i);
        newParent = parent.clone();
        reassign(0);
        for (int i = 1; i < n; i++) newChildren[newParent[i]].add(i);
        size = new int[n];
        measure(0);
        return size;
    }
    void reassign(int node) {
        char c = s.charAt(node);
        Deque<Integer> st = stacks.computeIfAbsent(c, k -> new ArrayDeque<>());
        if (!st.isEmpty()) newParent[node] = st.peek();
        st.push(node);
        for (int ch : children[node]) reassign(ch);
        st.pop();
    }
    int measure(int node) {
        int total = 1;
        for (int ch : newChildren[node]) total += measure(ch);
        size[node] = total;
        return total;
    }
}

class Solution {
public:
    vector<int> newParent, sz;
    vector<vector<int>> children, newChildren;
    unordered_map<char, vector<int>> stacks;
    string str;
    void reassign(int node) {
        char c = str[node];
        if (!stacks[c].empty()) newParent[node] = stacks[c].back();
        stacks[c].push_back(node);
        for (int ch : children[node]) reassign(ch);
        stacks[c].pop_back();
    }
    int measure(int node) {
        int total = 1;
        for (int ch : newChildren[node]) total += measure(ch);
        sz[node] = total;
        return total;
    }
    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        int n = parent.size(); str = s;
        children.assign(n, {}); newChildren.assign(n, {});
        for (int i = 1; i < n; i++) children[parent[i]].push_back(i);
        newParent = parent;
        reassign(0);
        for (int i = 1; i < n; i++) newChildren[newParent[i]].push_back(i);
        sz.assign(n, 0);
        measure(0);
        return sz;
    }
};

Explanation

Each node may move up to its closest ancestor sharing its letter. The closest such ancestor is simply the most recently entered ancestor with that letter as we walk down from the root, which a per-character stack captures perfectly.

The first DFS, reassign, walks the original tree. When we enter a node with letter c, the top of stacks[c] (if any) is exactly the deepest ancestor with letter c, so we set that as the node's new parent. Then we push the node onto stacks[c], recurse into its original children, and pop on the way back out so the stack always reflects the current root-to-node path.

All reattachments are computed from the original tree at once, so we record them into a newParent array rather than mutating as we go.

The second DFS, measure, rebuilds the children lists from newParent and computes each subtree size bottom-up: a node's size is one plus the sizes of its (possibly new) children.

Example on parent = [-1,0,0,1,1,1], s = "abaaba": nodes 3 and 5 (letter 'a') reattach to the root 0 (their closest 'a'-ancestor), node 4 (letter 'b') stays under node 1. Final sizes are [6,2,1,1,1,1].

Time: O(n) Space: O(n)