We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsCompaniesDisclaimersPrivacyTermsContactAbout Us

Find Subtree Sizes After Changes

Number: 3576

Difficulty: Medium

Paid? No

Companies: N/A

Problem Description

Given a tree with n nodes (rooted at node 0) described by an array parent and a string s representing the character assigned to each node, we perform one simultaneous update for every node (except the root). For each node x, we locate its closest ancestor y (if any) such that s[x] equals s[y]. If found, we remove the edge between x and its original parent and attach x as a child of y. Finally, we are required to compute and return an answer array where answer[i] is the size of the subtree rooted at node i in the resulting tree after the changes.

Key Insights

  • The update is done simultaneously; therefore, each node’s new parent is determined based solely on the ancestor relationship in the original tree.
  • A depth-first search (DFS) approach can traverse the tree and maintain a mapping (or stack) from character to the latest encountered node along the current DFS path.
  • For each node x (except the root), the "closest" ancestor with the same character (if exists) is given by the most recent entry in the mapping for s[x].
  • After determining the new parent for every node, we rebuild the tree (adjacency list) and perform a second DFS to compute the subtree sizes.

Space and Time Complexity

Time Complexity: O(n) – Each node is visited a constant number of times during the DFS traversals. Space Complexity: O(n) – Space is used for recursion stack, the mapping of last seen characters, and storing tree adjacencies.

Solution

We solve the problem in two major steps:

  1. Update Parent Relationships:

    • Build an adjacency list “children” from the original parent array.
    • Use a DFS starting at the root and maintain a dictionary (or map) “lastSeen” that maps a character to the most recent node (the closest ancestor encountered so far) with that character.
    • For each node (other than the root), before recursing on its children, determine the new parent by checking lastSeen for the node’s character.
    • After processing a node, backtrack the changes to lastSeen.

  2. Compute Subtree Sizes:

    • Rebuild the final tree using the newly computed parent array.
    • Run another DFS starting from the root to compute the subtree size for every node (subtree size equals 1 plus the sum of subtree sizes of all children).

This two-phase approach ensures that we correctly update the tree based on the original ancestry and then compute the required subtree sizes.

Code Solutions

Below are the code solutions for Python, JavaScript, C++, and Java with line-by-line comments.

from collections import defaultdict
import sys
sys.setrecursionlimit(10**6)

def findSubtreeSizes(parent, s):
    n = len(parent)
    # Build the original tree as an adjacency list
    children = defaultdict(list)
    for i in range(1, n):
        children[parent[i]].append(i)
    
    # newParent will hold the updated parent for each node
    newParent = [-1] * n
    newParent[0] = -1  # Root remains root

    # DFS helper function to update newParent with a mapping of last seen character
    lastSeen = {}  # Maps character to the last seen node index along current DFS path

    def dfs(node):
        # Save current letter and previous mapping value for backtracking
        char = s[node]
        prev = lastSeen.get(char, None)
        lastSeen[char] = node  # Update last seen for current character
        
        # Process each child of the current node
        for child in children[node]:
            # The candidate new parent for 'child' is the most recent occurrence 
            # of its character on the current DFS path.
            candidate = lastSeen.get(s[child], None)
            if candidate is None:
                # If no candidate exists, keep the original parent
                newParent[child] = parent[child]
            else:
                newParent[child] = candidate
            dfs(child)
        
        # Backtrack: restore lastSeen for current character
        if prev is None:
            del lastSeen[char]
        else:
            lastSeen[char] = prev

    dfs(0)

    # Build final tree based on newParent relationships
    finalChildren = [[] for _ in range(n)]
    for i in range(1, n):
        finalChildren[newParent[i]].append(i)

    # DFS to compute subtree sizes 
    answer = [0] * n
    def computeSize(node):
        count = 1
        for child in finalChildren[node]:
            count += computeSize(child)
        answer[node] = count
        return count

    computeSize(0)
    return answer

# Sample tests
if __name__ == "__main__":
    print(findSubtreeSizes([-1, 0, 0, 1, 1, 1], "abaabc"))  # Expected: [6, 3, 1, 1, 1, 1]
    print(findSubtreeSizes([-1, 0, 4, 0, 1], "abbba"))        # Expected: [5, 2, 1, 1, 1]
← Back to All Questions