Scale Engineer

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
ArrayHash TableStringTree

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 <= 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.

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.

  1. 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.

  2. Build New Tree: Using the newParent array, we construct an adjacency list representation of the final tree.

  3. 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.

import java.util.ArrayList; import java.util.List; class Solution {    // DFS to calculate subtree sizes    private int dfsSize(int u, List<List<Integer>> adj, int[] answer) {        int size = 1;        for (int v : adj.get(u)) {            size += dfsSize(v, adj, answer);        }        answer[u] = size;        return size;    }     public int[] countSubtrees(int n, int[] parent, String s) {        // Step 1: Determine new parents (Brute Force)        int[] newParent = parent.clone();        for (int i = 1; i < n; i++) {            int curr = parent[i];            while (curr != -1) {                if (s.charAt(curr) == s.charAt(i)) {                    newParent[i] = curr;                    break;                }                curr = parent[curr];            }        }         // Step 2: Build new tree        List<List<Integer>> newAdj = new ArrayList<>();        for (int i = 0; i < n; i++) {            newAdj.add(new ArrayList<>());        }        for (int i = 1; i < n; i++) {            if (newParent[i] != -1) {                 newAdj.get(newParent[i]).add(i);            }        }         // Step 3: Calculate subtree sizes        int[] answer = new int[n];        dfsSize(0, newAdj, answer);        return answer;    }}

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.

Solutions

class Solution {private  List<Integer>[] g;private  List<Integer>[] d;private  char[] s;private  int[] ans;public  int[] findSubtreeSizes(int[] parent, String s) {    int n = s.length();    g = new List[n];    d = new List[26];    this.s = s.toCharArray();    Arrays.setAll(g, k->new ArrayList<>());    Arrays.setAll(d, k->new ArrayList<>());    for (int i = 1; i < n; ++i) {      g[parent[i]].add(i);    }    ans = new int[n];    dfs(0, -1);    return ans;  }private  void dfs(int i, int fa) {    ans[i] = 1;    int idx = s[i] - 'a';    d[idx].add(i);    for (int j : g[i]) {      dfs(j, i);    }    int k = d[idx].size() > 1 ? d[idx].get(d[idx].size() - 2) : fa;    if (k >= 0) {      ans[k] += ans[i];    }    d[idx].remove(d[idx].size() - 1);  }}

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 →

Next problem