Find Subtree Sizes After Changes solution and explanation
5 more videos on this problem
Leetcode Biweekly Contest 142 | 3331. Find Subtree Sizes After Changes | CodefodCodeFod · 12:11 · 791 views
Leetcode 3331 | Find Subtree Sizes After Changes | Biweekly Contest 142Road To FAANG · 11:52 · 436 views
3331. Find Subtree Sizes After Changes (Leetcode Medium)Programming Live with Larry · 20:57 · 216 views
3331. Find Subtree Sizes After Changes | Tree | DFS | LeetCode | Medium | Biweekly Contest 142Leet's Code · 12:17 · 191 views
Biweekly Contest - 142 | First 3 Problems- 3330, 3331 & 3332 | LeetcodeAmit Choraria · 33:16 · 155 views
[English] 3331. Find Subtree Sizes After Changes Q2 Biweekly 142 [Leetcode Medium]Romain Lhotte · 14:40 · 107 views
The problem
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
yto nodexsuch thatyis an ancestor ofx, ands[x] == s[y]. - If node
ydoes not exist, do nothing. - Otherwise, remove the edge between
xand its current parent and make nodeythe new parent ofxby 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:
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:
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.length1 <= n <= 1050 <= parent[i] <= n - 1for alli >= 1.parent[0] == -1parentrepresents a valid tree.sconsists only of lowercase English letters.
Approach overview
Problem Overview: You are given a rooted tree defined by a parent array and a string s where s[i] is the label of node i. Each node may reconnect to the nearest ancestor that has the same character. After applying these changes simultaneously, compute the size of every node’s subtree in the new tree.
Approach 1: DFS with Ancestor Character Tracking (O(n) time, O(n) space)
Build the original tree using an adjacency list from the parent array, then run a Depth-First Search. During DFS, maintain a map from character to the most recent ancestor index with that character. When visiting node u, check whether its character already appears in the map. If it does, the new parent becomes that ancestor; otherwise the original parent remains. Maintain a temporary stack/list for each character so that when DFS backtracks the previous ancestor is restored. Once the new parent for each node is determined, construct the updated tree and run another DFS to compute subtree sizes. This works because every node is visited once and character lookups are constant-time via a hash table.
Approach 2: Union-Find Based Construction (O(n α(n)) time, O(n) space)
You can also treat the reattachment step as merging nodes with their closest matching ancestor using a Union-Find (Disjoint Set Union) structure. While traversing the tree with DFS, maintain the latest ancestor for each character. If node u finds a matching ancestor a, union u with a. The disjoint-set structure efficiently tracks component representatives with near-constant amortized cost α(n). After all unions are applied, rebuild the effective parent relationships and compute subtree sizes with a DFS on the resulting structure. This approach highlights connectivity management and works well when tree relationships are dynamically merged.
The core insight across both methods: every node only needs the closest ancestor with the same character. A DFS traversal naturally exposes the current ancestor chain, making it easy to track the latest node for each character. That eliminates expensive upward scans and keeps the solution linear.
Recommended for interviews: The DFS ancestor-tracking approach is the expected solution. It uses standard tree traversal, constant-time hash lookups, and processes each node exactly once for O(n) time. Discussing it shows strong understanding of DFS state management along the recursion path. The Union-Find version is a useful alternative when you want to reason about connectivity or demonstrate familiarity with disjoint-set structures.
Approach 1: Using DFS to Track Ancestors
This approach uses Depth First Search (DFS) to explore each node in the tree, simulating the effect of the changes by keeping track of a stack of ancestors' indices. When a node with the same character as the current node is found in the ancestor stack, we reassign the parent. Finally, we perform another DFS to calculate the subtree sizes accurately.
def findSubtreeSizes(parent, s):
n = len(parent)
ancestors = [[] for _ in range(n)]
for i in range(1, n):
ancestor = parent[i]
while ancestor != -1:
ancestors[i].append(ancestor)
ancestor = parent[ancestor]
new_parent = parent[:]
for i in range(n):
for ancestor in ancestors[i]:
if s[ancestor] == s[i]:
new_parent[i] = ancestor
break
children = [[] for _ in range(n)]
for i in range(1, n):
children[new_parent[i]].append(i)
def dfs(node):
subtree_size = 1
for child in children[node]:
subtree_size += dfs(child)
answer[node] = subtree_size
return subtree_size
answer = [0] * n
dfs(0)
return answer
# Example Usage:
print(findSubtreeSizes([-1,0,0,1,1,1], 'abaabc')) # Output: [6,3,1,1,1,1]
print(findSubtreeSizes([-1,0,4,0,1], 'abbba')) # Output: [5,2,1,1,1]import java.util.ArrayList;
import java.util.List;
public class TreeSubtreeSize {
public static int[] findSubtreeSizes(int[] parent, String s) {
int n = parent.length;
List<Integer>[] ancestors = new List[n];
for (int i = 0; i < n; i++) {
ancestors[i] = new ArrayList<>();
}
for (int i = 1; i < n; i++) {
int ancestor = parent[i];
while (ancestor != -1) {
ancestors[i].add(ancestor);
ancestor = parent[ancestor];
}
}
int[] newParent = parent.clone();
for (int i = 1; i < n; i++) {
for (int ancestor : ancestors[i]) {
if (s.charAt(ancestor) == s.charAt(i)) {
newParent[i] = ancestor;
break;
}
}
}
List<Integer>[] children = new List[n];
for (int i = 0; i < n; i++) {
children[i] = new ArrayList<>();
}
for (int i = 1; i < n; i++) {
children[newParent[i]].add(i);
}
int[] answer = new int[n];
dfs(0, children, answer);
return answer;
}
private static int dfs(int node, List<Integer>[] children, int[] answer) {
int subtreeSize = 1;
for (int child : children[node]) {
subtreeSize += dfs(child, children, answer);
}
answer[node] = subtreeSize;
return subtreeSize;
}
public static void main(String[] args) {
int[] result1 = findSubtreeSizes(new int[]{-1, 0, 0, 1, 1, 1}, "abaabc");
int[] result2 = findSubtreeSizes(new int[]{-1, 0, 4, 0, 1}, "abbba");
for (int num : result1) System.out.print(num + " ");
System.out.println();
for (int num : result2) System.out.print(num + " ");
}
}The algorithm first constructs a list of ancestors for each node, where each list contains all ancestors of the node in a top-down order. This is achieved by tracing back from the current node to the root following the given parent array.
Then, it checks for each node if there is an ancestor with the same character. If found, it updates the new parent for the node.
After updating the parent-child relationships for all nodes, a second DFS is performed to calculate the size of the subtree rooted at each node. This is accomplished by summing up sizes of subtrees of its children and itself.
Time Complexity: O(n), because each node and its ancestors are visited a constant number of times.
Space Complexity: O(n), due to additional space for the ancestor tracking and recursion stack in DFS.
Approach 2: Using Union-Find Data Structure
This alternative approach applies a Union-Find data structure, also known as Disjoint Set Union (DSU), to dynamically manage the reparenting of nodes. We utilize the union-find structure to merge nodes with the same character whenever changes are detected.
class UnionFind:
def __init__(self, size):
self.parent = list(range(size))
self.rank = [0] * size
def find(self, x):
if self.parent[x] != x:
self.parent[x] = self.find(self.parent[x])
return self.parent[x]
def union(self, x, y):
rootX = self.find(x)
rootY = self.find(y)
if rootX != rootY:
if self.rank[rootX] > self.rank[rootY]:
self.parent[rootY] = rootX
elif self.rank[rootX] < self.rank[rootY]:
self.parent[rootX] = rootY
else:
self.parent[rootY] = rootX
self.rank[rootX] += 1
def findSubtreeSizes(parent, s):
n = len(parent)
uf = UnionFind(n)
char_to_index = {}
for i in range(n):
if s[i] in char_to_index:
uf.union(i, char_to_index[s[i]])
char_to_index[s[i]] = i
root_to_size = {}
for i in range(n):
root = uf.find(i)
if root in root_to_size:
root_to_size[root] += 1
else:
root_to_size[root] = 1
answer = [0] * n
for i in range(n):
answer[i] = root_to_size[uf.find(i)]
return answer
# Example Usage:
print(findSubtreeSizes([-1,0,0,1,1,1], 'abaabc')) # Output: [6,3,1,1,1,1]
print(findSubtreeSizes([-1,0,4,0,1], 'abbba')) # Output: [5,2,1,1,1]This solution initializes a Union-Find (or Disjoint Set Union) structure to manage node grouping.
As each node is processed, if another node with the same character has been encountered, it utilizes union operations to connect these nodes. After this step, all nodes sharing the same character and in an ancestor relationship are labeled with a unique group identifier.
Then, by simply counting the group sizes, the size of the final subtrees can be obtained for each node.
Time Complexity: O(n * α(n)), where α(n) is the inverse Ackermann function, a very slowly growing function. Effectively O(n) for this problem's constraints.
Space Complexity: O(n), used for Union-Find data structure and size tracking.
Approach 3: Default Approach
class Solution {
public:
vector<int> findSubtreeSizes(vector<int>& parent, string s) {
int n = s.size();
vector<int> g[n];
vector<int> d[26];
for (int i = 1; i < n; ++i) {
g[parent[i]].push_back(i);
}
vector<int> ans(n);
auto dfs = [&](this auto&& dfs, int i, int fa) -> void {
ans[i] = 1;
int idx = s[i] - 'a';
d[idx].push_back(i);
for (int j : g[i]) {
dfs(j, i);
}
int k = d[idx].size() > 1 ? d[idx][d[idx].size() - 2] : fa;
if (k >= 0) {
ans[k] += ans[i];
}
d[idx].pop_back();
};
dfs(0, -1);
return ans;
}
};function findSubtreeSizes(parent: number[], s: string): number[] {
const n = parent.length;
const g: number[][] = Array.from({ length: n }, () => []);
const d: number[][] = Array.from({ length: 26 }, () => []);
for (let i = 1; i < n; ++i) {
g[parent[i]].push(i);
}
const ans: number[] = Array(n).fill(1);
const dfs = (i: number, fa: number): void => {
const idx = s.charCodeAt(i) - 97;
d[idx].push(i);
for (const j of g[i]) {
dfs(j, i);
}
const k = d[idx].length > 1 ? d[idx].at(-2)! : fa;
if (k >= 0) {
ans[k] += ans[i];
}
d[idx].pop();
};
dfs(0, -1);
return ans;
}func findSubtreeSizes(parent []int, s string) []int {
n := len(s)
g := make([][]int, n)
for i := 1; i < n; i++ {
g[parent[i]] = append(g[parent[i]], i)
}
d := [26][]int{}
ans := make([]int, n)
var dfs func(int, int)
dfs = func(i, fa int) {
ans[i] = 1
idx := int(s[i] - 'a')
d[idx] = append(d[idx], i)
for _, j := range g[i] {
dfs(j, i)
}
k := fa
if len(d[idx]) > 1 {
k = d[idx][len(d[idx])-2]
}
if k != -1 {
ans[k] += ans[i]
}
d[idx] = d[idx][:len(d[idx])-1]
}
dfs(0, -1)
return ans
}Complexity comparison
| Approach | Complexity |
|---|---|
| Using DFS to Track Ancestors | Time Complexity: O(n), because each node and its ancestors are visited a constant number of times. Space Complexity: O(n), due to additional space for the ancestor tracking and recursion stack in DFS. |
| Using Union-Find Data Structure | Time Complexity: O(n * α(n)), where α(n) is the inverse Ackermann function, a very slowly growing function. Effectively O(n) for this problem's constraints. Space Complexity: O(n), used for Union-Find data structure and size tracking. |
| Default Approach | – |
Complexity analysis
| Approach | Time | Space | When to Use |
|---|---|---|---|
| DFS with Ancestor Character Tracking | O(n) | O(n) | Best general solution. Efficient when processing ancestor relationships during a DFS traversal. |
| Union-Find (Disjoint Set) | O(n α(n)) | O(n) | Useful when modeling node reattachments as connectivity merges or demonstrating DSU knowledge. |
Frequently asked questions
Is Find Subtree Sizes After Changes easy or hard?
Find Subtree Sizes After Changes is rated Medium difficulty. The challenge comes from recognizing that the nearest matching ancestor can be tracked during DFS instead of scanning upward in the tree. Once that insight is applied, the rest of the problem becomes a standard subtree size computation.
Find Subtree Sizes After Changes Python/Java solution
Python and Java solutions typically implement DFS with a dictionary or map that records the latest node for each character along the current path. The traversal updates parent relationships and later computes subtree sizes using another DFS. Both implementations run in O(n) time and require O(n) additional memory.
How to solve Find Subtree Sizes After Changes in O(n)?
Perform a DFS from the root while maintaining a map of the latest ancestor index for each character. For each node, check if the same character already exists on the current path and redirect its parent to that ancestor. After determining all updated parent relationships, build the new tree and run another DFS to compute subtree sizes. Every node is processed a constant number of times.
What is the best approach for Find Subtree Sizes After Changes?
The most efficient approach uses DFS while tracking the latest ancestor for each character. During traversal, maintain a hash map from character to the most recent node index on the current DFS path. This lets you find the closest matching ancestor in O(1) time and update the effective parent. The full algorithm runs in O(n) time and O(n) space.
Is Find Subtree Sizes After Changes asked at Google/Amazon/Meta?
Tree traversal and ancestor tracking problems like this appear frequently in interviews at companies such as Google, Amazon, and Meta. Variants often involve DFS state tracking, nearest ancestor queries, or dynamic tree restructuring. Understanding how to maintain path information during DFS is a commonly tested skill.
What data structure is used in Find Subtree Sizes After Changes?
The main structures are adjacency lists for representing the tree, a hash table for tracking the latest ancestor of each character, and recursion or an explicit stack for DFS traversal. Some alternative solutions also use the Union-Find data structure to manage node merges efficiently.
What is the time complexity of Find Subtree Sizes After Changes?
The optimal solution runs in O(n) time where n is the number of nodes in the tree. Each node is visited once during DFS, and ancestor lookups are constant-time using a hash map. Computing subtree sizes also requires a single traversal, keeping the overall complexity linear.