3331. Find Subtree Sizes After Changes LeetCode Solution

Last updated: March 6, 2025

In this guide, you will get 3331. Find Subtree Sizes After Changes LeetCode Solution with the best time and space complexity. The solution to Find Subtree Sizes After Changes problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.

Problem Statement
Complexity Analysis
Find Subtree Sizes After Changes solution in C++
Find Subtree Sizes After Changes solution in Java
Find Subtree Sizes After Changes solution in Python
Additional Resources

3331. Find Subtree Sizes After Changes LeetCode Solution image

Problem Statement of Find Subtree Sizes After Changes

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] = 1.
parent[0] == -1
parent represents a valid tree.
s consists only of lowercase English letters.

Complexity Analysis

Time Complexity: O(nh)
Space Complexity: O(n)

3331. Find Subtree Sizes After Changes LeetCode Solution in C++

class Solution {
 public:
  vector<int> findSubtreeSizes(vector<int>& parent, string s) {
    const int n = parent.size();
    vector<int> ans(n);
    vector<int> newParent = parent;
    vector<vector<int>> tree(n);

    for (int i = 1; i < n; ++i) {
      const int closest = findClosestAncestor(i, parent, s);
      if (closest != -1)
        newParent[i] = closest;
    }

    for (int i = 1; i < n; ++i)
      tree[newParent[i]].push_back(i);

    dfs(tree, 0, ans);
    return ans;
  }

 private:
  // Returns the closest ancestor of node `u` that has the same value as `u`.
  int findClosestAncestor(int u, const vector<int>& parent, const string& s) {
    for (int curr = parent[u]; curr != -1; curr = parent[curr])
      if (s[curr] == s[u])
        return curr;
    return -1;
  }

  int dfs(const vector<vector<int>>& tree, int u, vector<int>& ans) {
    int sz = 1;
    for (const int v : tree[u])
      sz += dfs(tree, v, ans);
    return ans[u] = sz;
  }
};
/* code provided by PROGIEZ */

3331. Find Subtree Sizes After Changes LeetCode Solution in Java

class Solution {
  public int[] findSubtreeSizes(int[] parent, String s) {
    final int n = parent.length;
    int[] ans = new int[n];
    int[] newParent = parent.clone();
    List<Integer>[] tree = new ArrayList[n];

    for (int i = 0; i < n; ++i)
      tree[i] = new ArrayList<>();

    for (int i = 1; i < n; ++i) {
      final int closest = findClosestAncestor(i, parent, s);
      if (closest != -1)
        newParent[i] = closest;
    }

    for (int i = 1; i < n; ++i)
      tree[newParent[i]].add(i);

    dfs(tree, 0, ans);
    return ans;
  }

  // Returns the closest ancestor of node `u` that has the same value as `u`.
  private int findClosestAncestor(int u, int[] parent, String s) {
    for (int curr = parent[u]; curr != -1; curr = parent[curr])
      if (s.charAt(curr) == s.charAt(u))
        return curr;
    return -1;
  }

  private int dfs(List<Integer>[] tree, int u, int[] ans) {
    int sz = 1;
    for (final int v : tree[u])
      sz += dfs(tree, v, ans);
    return ans[u] = sz;
  }
}
// code provided by PROGIEZ

3331. Find Subtree Sizes After Changes LeetCode Solution in Python

class Solution:
  def findSubtreeSizes(self, parent: list[int], s: str) -> list[int]:
    n = len(parent)
    ans = [0] * n
    newParent = parent.copy()
    tree = [[] for _ in range(n)]

    for i in range(1, n):
      closest = self._findClosestAncestor(i, parent, s)
      if closest != -1:
        newParent[i] = closest

    for i in range(1, n):
      tree[newParent[i]].append(i)

    self._dfs(tree, 0, ans)
    return ans

  def _findClosestAncestor(self, u: int, parent: list[int], s: str) -> int:
    """
    Returns the closest ancestor of node `u` that has the same value as `u`.
    """
    curr = parent[u]
    while curr != -1:
      if s[curr] == s[u]:
        return curr
      curr = parent[curr]
    return -1

  def _dfs(self, tree: list[list[int]], u: int, ans: list[int]) -> int:
    sz = 1
    for v in tree[u]:
      sz += self._dfs(tree, v, ans)
    ans[u] = sz
    return sz
 # code by PROGIEZ

Additional Resources

Happy Coding! Keep following PROGIEZ for more updates and solutions.

Table of Contents