Find Subtree Sizes After Changes

Med

#2958Time: O(n^2) - The dominant part is finding the new parents. For each of the `n` nodes, we might traverse up to `n` ancestors in the worst case (a skewed tree), leading to a quadratic time complexity.Space: O(n) - To store the `newParent` array, the adjacency list for the new tree (`newAdj`), and the recursion stack for the DFS, all of which can be proportional to `n`.

Algorithms

Depth-First Search

Data structures

Array Hash Table String Tree

Prompt

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:

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.length
1 <= n <= 10⁵
0 <= parent[i] <= n - 1 for all i >= 1.
parent[0] == -1
parent represents a valid tree.
s consists only of lowercase English letters.

Approaches

2 approaches with complexity analysis and trade-offs.

This approach directly simulates the process described in the problem. For each node, it finds its new parent by naively traversing up its chain of ancestors in the original tree. Once all new parent relationships are determined, it builds the new tree and then runs a Depth-First Search (DFS) to calculate the size of each subtree.

Algorithm

Step 1: Determine New Parents (Brute-Force)
- Create a newParent array, initialized as a copy of the input parent array.
- Iterate through each node x from 1 to n-1.
- For each x, start a traversal from its current parent p = parent[x].
- In a while loop, traverse up the tree towards the root (p != -1).
- If the character of the ancestor p matches the character of x (s[p] == s[x]), we have found the closest ancestor. Update newParent[x] = p and stop searching for this x.
- If the characters don't match, move to the next ancestor: p = parent[p].
Step 2: Build the Final Tree
- After computing all new parent relationships, construct the final tree structure.
- Create an adjacency list, newAdj, to represent the new tree.
- Iterate from i = 1 to n-1 and for each i, add a directed edge from newParent[i] to i.
Step 3: Calculate Subtree Sizes
- Perform a Depth-First Search (DFS) on the newly constructed tree, starting from the root (node 0).
- The DFS function, dfs_size(u), will recursively calculate the size of the subtree rooted at u.
- The size of a subtree at u is 1 (for u itself) plus the sum of the sizes of the subtrees of its children.
- Store the computed sizes in a result array.

Walkthrough

The core of this method is a nested loop structure to find the new parents. The outer loop iterates through each node, and the inner loop traverses its ancestors. This is straightforward but inefficient.

Find New Parents: We iterate through every node i from 1 to n-1. For each node, we walk up from its parent, parent[i], to the root, checking each ancestor. The first ancestor found with a matching character becomes the new parent. This is stored in a newParent array.
Build New Tree: Using the newParent array, we construct an adjacency list representation of the final tree.
Calculate Sizes: A standard DFS traversal on the new tree calculates the subtree sizes. The size of a node's subtree is 1 plus the sum of its children's subtree sizes.

1import java.util.ArrayList; 2import java.util.List;3 4class Solution {5    // DFS to calculate subtree sizes6    private int dfsSize(int u, List<List<Integer>> adj, int[] answer) {7        int size = 1;8        for (int v : adj.get(u)) {9            size += dfsSize(v, adj, answer);10        }11        answer[u] = size;12        return size;13    }14 15    public int[] countSubtrees(int n, int[] parent, String s) {16        // Step 1: Determine new parents (Brute Force)17        int[] newParent = parent.clone();18        for (int i = 1; i < n; i++) {19            int curr = parent[i];20            while (curr != -1) {21                if (s.charAt(curr) == s.charAt(i)) {22                    newParent[i] = curr;23                    break;24                }25                curr = parent[curr];26            }27        }28 29        // Step 2: Build new tree30        List<List<Integer>> newAdj = new ArrayList<>();31        for (int i = 0; i < n; i++) {32            newAdj.add(new ArrayList<>());33        }34        for (int i = 1; i < n; i++) {35            if (newParent[i] != -1) {36                 newAdj.get(newParent[i]).add(i);37            }38        }39 40        // Step 3: Calculate subtree sizes41        int[] answer = new int[n];42        dfsSize(0, newAdj, answer);43        return answer;44    }45}

Complexity

Time

O(n^2) - The dominant part is finding the new parents. For each of the `n` nodes, we might traverse up to `n` ancestors in the worst case (a skewed tree), leading to a quadratic time complexity.

Space

O(n) - To store the `newParent` array, the adjacency list for the new tree (`newAdj`), and the recursion stack for the DFS, all of which can be proportional to `n`.

Trade-offs

Pros

Simple to understand and implement as it directly follows the problem description.

Cons

The time complexity of O(n^2) is too slow for the given constraints (n <= 10^5) and will result in a 'Time Limit Exceeded' error.

This optimized approach avoids the expensive re-computation of ancestor paths by using a single Depth-First Search (DFS) on the original tree. During this traversal, it cleverly keeps track of the most recent ancestor for each character along the current path. This allows finding the new parent for each node in O(1) time (amortized). After all new parents are found, a second DFS is performed on the final tree to calculate subtree sizes.

Algorithm

Step 1: Build Original Tree and Initialize
- Construct an adjacency list originalAdj for the original tree from the parent array to enable efficient traversal.
- Initialize a newParent array, copying the original parent values.
- Create a helper array lastSeen[26] initialized to a sentinel value (e.g., -1). This array will track the most recent ancestor for each character on the current DFS path.
Step 2: Find New Parents with a Single DFS
- Perform a single DFS traversal on the original tree, starting from the root (node 0).
- In the dfs(u) function:
  - Find the closest ancestor for u with character s[u] by checking lastSeen[s.charAt(u) - 'a']. If an ancestor exists, update newParent[u].
  - Before visiting the children of u, update the lastSeen state: save the old value for s[u]'s character and set lastSeen[s.charAt(u) - 'a'] = u.
  - Recursively call the DFS for all children of u.
  - After the recursive calls return (after visiting u's entire subtree), backtrack by restoring the lastSeen array to its previous state. This is critical for correctness.
Step 3: Build New Tree and Calculate Sizes
- This part is the same as the brute-force approach.
- Build a new adjacency list newAdj from the final newParent array.
- Run a second DFS on newAdj to compute the subtree sizes for all nodes.

Walkthrough

The key optimization is to avoid the O(n) ancestor scan for each node. By traversing the original tree with DFS, we can maintain the path from the root to the current node. A helper array, lastSeen, indexed by character, stores the deepest node on the current path for each character. When we visit a node u, lastSeen[s.charAt(u)] gives us its closest ancestor with the same character instantly.

Crucially, we must use backtracking. When the DFS for a subtree rooted at u is complete, we must restore the lastSeen array to the state it was in before visiting u. This ensures that when the traversal moves to a different branch (e.g., a sibling of u), the lastSeen information is correct for that new path.

After this O(n) pass determines all newParent relationships, we build the new tree and run a final O(n) DFS to find the subtree sizes, just like in the previous approach.

1import java.util.ArrayList;2import java.util.Arrays;3import java.util.List;4 5class Solution {6    public int[] countSubtrees(int n, int[] parent, String s) {7        // Step 1: Build original tree's adjacency list8        List<List<Integer>> originalAdj = new ArrayList<>();9        for (int i = 0; i < n; i++) {10            originalAdj.add(new ArrayList<>());11        }12        for (int i = 1; i < n; i++) {13            originalAdj.get(parent[i]).add(i);14        }15 16        // Step 2: Find new parents using an optimized DFS17        int[] newParent = parent.clone();18        int[] lastSeen = new int[26];19        Arrays.fill(lastSeen, -1);20        findNewParentsDfs(0, originalAdj, s, lastSeen, newParent);21 22        // Step 3: Build the new tree's adjacency list23        List<List<Integer>> newAdj = new ArrayList<>();24        for (int i = 0; i < n; i++) {25            newAdj.add(new ArrayList<>());26        }27        for (int i = 1; i < n; i++) {28            if (newParent[i] != -1) {29                newAdj.get(newParent[i]).add(i);30            }31        }32 33        // Step 4: Calculate subtree sizes on the new tree34        int[] answer = new int[n];35        calculateSizesDfs(0, newAdj, answer);36        return answer;37    }38 39    private void findNewParentsDfs(int u, List<List<Integer>> adj, String s, int[] lastSeen, int[] newParent) {40        int charIndex = s.charAt(u) - 'a';41 42        // Find closest ancestor with the same character from the current path43        int ancestor = lastSeen[charIndex];44        if (ancestor != -1) {45            newParent[u] = ancestor;46        }47 48        // Update lastSeen for descendants and recurse49        int prevSeenNode = lastSeen[charIndex];50        lastSeen[charIndex] = u;51 52        for (int v : adj.get(u)) {53            findNewParentsDfs(v, adj, s, lastSeen, newParent);54        }55 56        // Backtrack: restore lastSeen to its previous state for other branches57        lastSeen[charIndex] = prevSeenNode;58    }59 60    private int calculateSizesDfs(int u, List<List<Integer>> adj, int[] answer) {61        int size = 1;62        for (int v : adj.get(u)) {63            size += calculateSizesDfs(v, adj, answer);64        }65        answer[u] = size;66        return size;67    }68}

Complexity

Time

O(n) - The solution involves a few linear passes over the data. Building the original tree is O(n). The first DFS to find parents is O(n). Building the new tree is O(n). The second DFS to find sizes is O(n). The total complexity is O(n).

Space

O(n) - For storing two adjacency lists, `newParent` array, `answer` array, and the recursion stacks for two DFS traversals. The depth of the recursion stack can be up to O(n) in a skewed tree.

Trade-offs

Pros

Optimal time complexity of O(n), which is efficient enough for the given constraints.
The logic is clean and separates the problem into two distinct phases: determining the new structure and analyzing it.

Cons

More complex to implement due to the two-pass DFS structure and the need for careful state management (backtracking) in the first DFS.

Solutions

Solution

1class Solution {2private3  List<Integer>[] g;4private5  List<Integer>[] d;6private7  char[] s;8private9  int[] ans;10public11  int[] findSubtreeSizes(int[] parent, String s) {12    int n = s.length();13    g = new List[n];14    d = new List[26];15    this.s = s.toCharArray();16    Arrays.setAll(g, k->new ArrayList<>());17    Arrays.setAll(d, k->new ArrayList<>());18    for (int i = 1; i < n; ++i) {19      g[parent[i]].add(i);20    }21    ans = new int[n];22    dfs(0, -1);23    return ans;24  }25private26  void dfs(int i, int fa) {27    ans[i] = 1;28    int idx = s[i] - 'a';29    d[idx].add(i);30    for (int j : g[i]) {31      dfs(j, i);32    }33    int k = d[idx].size() > 1 ? d[idx].get(d[idx].size() - 2) : fa;34    if (k >= 0) {35      ans[k] += ans[i];36    }37    d[idx].remove(d[idx].size() - 1);38  }39}

1class Solution {2public:3  vector<int> findSubtreeSizes(vector<int> &parent, string s) {4    int n = s.size();5    vector<int> g[n];6    vector<int> d[26];7    for (int i = 1; i < n; ++i) {8      g[parent[i]].push_back(i);9    }10    vector<int> ans(n);11    auto dfs = [&](auto &&dfs, int i, int fa) -> void {12      ans[i] = 1;13      int idx = s[i] - 'a';14      d[idx].push_back(i);15      for (int j : g[i]) {16        dfs(dfs, j, i);17      }18      int k = d[idx].size() > 1 ? d[idx][d[idx].size() - 2] : fa;19      if (k >= 0) {20        ans[k] += ans[i];21      }22      d[idx].pop_back();23    };24    dfs(dfs, 0, -1);25    return ans;26  }27};

1class Solution:2    def findSubtreeSizes(self, parent: List[int], s: str) -> List[int]: def dfs(i: int, fa: int): ans[i] = 1 d[s[i]]. append(i) for j in g[i]: dfs(j, i) k = fa if len(d[s[i]]) > 1: k = d[s[i]][- 2] if k != - 1: ans[k] += ans[i] d[s[i]]. pop() n = len(s) g = [[] for _ in range(n)] for i in range(1, n): g[parent[i]]. append(i) d = defaultdict(list) ans = [0] * n dfs(0, - 1) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem