# Find Subtree Sizes After Changes
**Difficulty:** MEDIUM
[External](https://leetcode.com/problems/find-subtree-sizes-after-changes)
Canonical: https://scaleengineer.com/dsa/problems/find-subtree-sizes-after-changes
**Algorithms:** [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search)
**Data structures:** Array, Hash Table, String, Tree
---
## 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 `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.glich.co/dsa/find-subtree-sizes-after-changes/image0.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.glich.co/dsa/find-subtree-sizes-after-changes/image1.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.

# Approaches
## Brute-Force Simulation
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.
**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`.
**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.
### Explanation
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.

```java
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;
    }
}
```
### 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.

## Optimized Parent Finding with DFS
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.
**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.
**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.
### Explanation
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.

```java
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

class Solution {
    public int[] countSubtrees(int n, int[] parent, String s) {
        // Step 1: Build original tree's adjacency list
        List<List<Integer>> originalAdj = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            originalAdj.add(new ArrayList<>());
        }
        for (int i = 1; i < n; i++) {
            originalAdj.get(parent[i]).add(i);
        }

        // Step 2: Find new parents using an optimized DFS
        int[] newParent = parent.clone();
        int[] lastSeen = new int[26];
        Arrays.fill(lastSeen, -1);
        findNewParentsDfs(0, originalAdj, s, lastSeen, newParent);

        // Step 3: Build the new tree's adjacency list
        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 4: Calculate subtree sizes on the new tree
        int[] answer = new int[n];
        calculateSizesDfs(0, newAdj, answer);
        return answer;
    }

    private void findNewParentsDfs(int u, List<List<Integer>> adj, String s, int[] lastSeen, int[] newParent) {
        int charIndex = s.charAt(u) - 'a';

        // Find closest ancestor with the same character from the current path
        int ancestor = lastSeen[charIndex];
        if (ancestor != -1) {
            newParent[u] = ancestor;
        }

        // Update lastSeen for descendants and recurse
        int prevSeenNode = lastSeen[charIndex];
        lastSeen[charIndex] = u;

        for (int v : adj.get(u)) {
            findNewParentsDfs(v, adj, s, lastSeen, newParent);
        }

        // Backtrack: restore lastSeen to its previous state for other branches
        lastSeen[charIndex] = prevSeenNode;
    }

    private int calculateSizesDfs(int u, List<List<Integer>> adj, int[] answer) {
        int size = 1;
        for (int v : adj.get(u)) {
            size += calculateSizesDfs(v, adj, answer);
        }
        answer[u] = size;
        return size;
    }
}
```
### 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.

# Solutions
### Java

```java
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);
  }
}

```

### CPP

```cpp
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 = [&](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(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(dfs, 0, -1);
    return ans;
  }
};

```

### Python

```python
class Solution:
    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

```