Find Subtree Sizes After Changes
Source: leetcode
Difficulty: Medium
Tags: array, hash-table, string, tree, depth-first-search

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.


 


**Example 1:**



**Input:** parent \= \[\-1,0,0,1,1,1], s \= "abaabc"


**Output:** \[6,3,1,1,1,1]


**Explanation:**


![](https://assets.leetcode.com/uploads/2024/08/15/graphex1drawio.png)
The parent of node 3 will change from node 1 to node 0\.



**Example 2:**



**Input:** parent \= \[\-1,0,4,0,1], s \= "abbba"


**Output:** \[5,2,1,1,1]


**Explanation:**


![](https://assets.leetcode.com/uploads/2024/08/20/exgraph2drawio.png)
The following changes will happen at the same time:


- The parent of node 4 will change from node 1 to node 0\.
- The parent of node 2 will change from node 4 to node 1\.



 


**Constraints:**


- `n == parent.length == s.length`
- `1 <= n <= 105`
- `0 <= parent[i] <= n - 1` for all `i >= 1`.
- `parent[0] == -1`
- `parent` represents a valid tree.
- `s` consists only of lowercase English letters.




C++ solution:
class Solution {
public:
    // we're storing the occurences of each char in a vec //
    // size: (26,_) - one vec for each char seen so far
    vector<vector<int>> state;  

    vector<vector<int>> adj;
    vector<bool> vis;
    vector<int> subtreeSizes;
    
    void dfs(int node, string& s) {
        vis[node]=true;
        subtreeSizes[node] = 1; 

        state[s[node]-'a'].push_back(node);
        for(auto v: adj[node]) {
            if(!vis[v]) {
                dfs(v, s);
                /** above dfs will calculate subtree size rooted at v, now the imp part is which ancestor to add this zie to?? */
                
                // if s[v]-'a' has a matching, add to matched node
                int ind = s[v]-'a';
                if(state[ind].empty()) subtreeSizes[node] += subtreeSizes[v]; 
                else subtreeSizes[state[ind].back()] += subtreeSizes[v]; 
            }
        }
        state[s[node]-'a'].pop_back();     // backtrack
    }

    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        int n = parent.size();
        adj.resize(n); 
        for(int i=1; i<n; i++) {
            adj[parent[i]].push_back(i);
        }

        state.resize(26);
        vis.assign(n, false);
        subtreeSizes.resize(n, 0);
        dfs(0, s);
        return subtreeSizes;
    }
};

Go solution:
type Tree struct {
	Value    byte
	Node     int
	Children []*Tree
}

func findSubtreeSizes(parent []int, s string) []int {
	root := buildTree(parent, s)

	hash := make(map[byte]int)
	dfs(root, hash, parent)

	res := make([]int, len(parent))
	root2 := buildTree(parent, s)

	getSubTSize(root2, res)
	return res
}

func dfs(root *Tree, hash map[byte]int, parent []int) {
	if root == nil {
		return
	}
	tmp, exist := hash[root.Value]
	hash[root.Value] = root.Node
	for _, child := range root.Children {
		dfs(child, hash, parent)
	}
	if exist {
		parent[root.Node] = tmp
		hash[root.Value] = tmp
	} else {
		delete(hash, root.Value)
	}
}

func getSubTSize(root *Tree, res []int) int {
	if root == nil {
		return 0
	}
	ans := 1
	for _, child := range root.Children {
		ans += getSubTSize(child, res)
	}
	res[root.Node] = ans
	return ans
}

func buildTree(parent []int, s string) *Tree {
	n := len(parent)
	nodes := make([]*Tree, n)

	for i := 0; i < n; i++ {
		nodes[i] = &Tree{
			Value: s[i],
			Node:  i,
		}
	}

	root := nodes[0]
	for i := 1; i < n; i++ {
		parentNode := nodes[parent[i]]
		parentNode.Children = append(parentNode.Children, nodes[i])
	}

	return root
}

Java solution:
class Solution {
    public int[] findSubtreeSizes(int[] parent, String s) {
        int n = parent.length;
        List<List<Integer>> adj = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            adj.add(new ArrayList<>());
        }

        for (int i = 1; i < n; i++) {
            adj.get(parent[i]).add(i);
        }

        int[] newParent = Arrays.copyOf(parent, n);
        Map<Character, Integer> lastSeen = new HashMap<>();
        dfsAdjustParent(0, s, adj, lastSeen, newParent);

        List<List<Integer>> finalAdj = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            finalAdj.add(new ArrayList<>());
        }
        for (int i = 1; i < n; i++) {
            finalAdj.get(newParent[i]).add(i);
        }

        int[] answer = new int[n];
        dfsCountSubtreeSizes(0, finalAdj, answer);
        
        return answer;
    }

    public void dfsAdjustParent(int node, String s, List<List<Integer>> adj, Map<Character, Integer> lastSeen, int[] newParent) {
        char c = s.charAt(node);
        Integer prev = lastSeen.get(c);
        if (prev != null && prev != newParent[node]) {
            newParent[node] = prev;
        }

        Integer originalParent = lastSeen.put(c, node);
        for (int child : adj.get(node)) {
            dfsAdjustParent(child, s, adj, lastSeen, newParent);
        }
        if (originalParent == null) {
            lastSeen.remove(c);
        } else {
            lastSeen.put(c, originalParent);
        }
    }

    public int dfsCountSubtreeSizes(int node, List<List<Integer>> adj, int[] answer) {
        int size = 1;
        for (int child : adj.get(node)) {
            size += dfsCountSubtreeSizes(child, adj, answer);
        }
        answer[node] = size;
        return size;
    }
}

JavaScript solution:
/**
 * @param {number[]} parent
 * @param {string} s
 * @return {number[]}
 */
var findSubtreeSizes = function(parent, s) {
    const n = parent.length;
    const arr = Array.from({ length: n }, () => new Array());

    for (let i = 1; i < n; ++i) {
        arr[parent[i]].push(i);
    }

    const dfs = (node, map) => {
        const prev = map.get(s[node]);
        if (prev !== undefined) parent[node] = prev;
        map.set(s[node], node);
        for (let c of arr[node]) {
            dfs(c, map);
        }
        map.set(s[node], prev);
    };

    dfs(0, new Map());

    const result = new Array(n).fill(0);
    const arrSum = Array.from({ length: n }, () => new Array());

    for (let i = 1; i < n; ++i) {
        arrSum[parent[i]].push(i);
    }

    const calc = node => {
        let sum = 1;
        for (let c of arrSum[node]) {
            sum += calc(c);
        }
        return result[node] += sum;
    };

    calc(0);

    return result;
};

Python solution:
class Solution:
    def findSubtreeSizes(self, p: List[int], s: str) -> List[int]:
        n, clds = len(s), defaultdict(set)
        for i in range(n):
            clds[p[i]].add(i)
        res = [0] * n

        def dfs(i, anc={}):
            prv, anc[s[i]] = anc.get(s[i]), i
            res[i] = cnt = 1
            for c in clds[i]:
                dfs(c, anc)
                res[p[c]] += res[c]
            anc[s[i]] = prv
            if prv is not None:
                p[i] = prv

        dfs(0)
        return res

Ruby solution:
def find_subtree_sizes(parent, s)
  n = parent.size
  children_orig = Array.new(n) { [] }
  (1...n).each { |i| children_orig[parent[i]] << i }
  new_parent = Array.new(n)
  new_parent[0] = -1
  last = Array.new(26)
  stack = []
  li = s[0].ord - 97
  prev = last[li]
  last[li] = 0
  stack.push([0, 0, li, prev])
  while !stack.empty?
    node, child_idx, li_node, prev_val = stack.last
    if child_idx < children_orig[node].size
      child = children_orig[node][child_idx]
      stack.last[1] += 1
      li_child = s[child].ord - 97
      candidate = last[li_child]
      new_parent[child] = candidate.nil? ? parent[child] : candidate
      old = last[li_child]
      last[li_child] = child
      stack.push([child, 0, li_child, old])
    else
      frame = stack.pop
      last[frame[2]] = frame[3]
    end
  end
  children_final = Array.new(n) { [] }
  (1...n).each { |i| children_final[new_parent[i]] << i if new_parent[i] != -1 }
  ans = Array.new(n, 1)
  order = []
  stack = [0]
  while !stack.empty?
    cur = stack.pop
    order << cur
    children_final[cur].each { |child| stack.push(child) }
  end
  order.reverse_each { |node| ans[new_parent[node]] += ans[node] if new_parent[node] != -1 }
  ans
end