Find Subtree Sizes After Changes - Practice Coding Problems

Find Subtree Sizes After Changes - Problem

Array Hash Table String Tree Depth-First Search Medium

Imagine you have a family tree where each person has a unique character assigned to them. Now, here's the twist: every person wants to be closer to their ancestor who shares the same character!

You're given a tree rooted at node 0 with n nodes numbered from 0 to n-1. The tree structure is represented by a parent array where parent[i] is the parent of node i (with parent[0] = -1 since node 0 is the root). Each node i has a character s[i] assigned to it.

Here's what happens simultaneously for all nodes from 1 to n-1:

Each node x looks for the closest ancestor y that has the same character (s[x] == s[y])
If such an ancestor exists, node x abandons its current parent and becomes a direct child of ancestor y
If no such ancestor exists, the node stays put

Your task is to return an array where answer[i] represents the size of the subtree rooted at node i after all these family relocations are complete!

Input & Output

example_1.py — Basic Tree Transformation

$ Input: parent = [-1,0,1,2], s = "abbc"

› Output: [4,2,1,1]

💡 Note: Node 2 (character 'b') finds ancestor node 1 (also 'b') and becomes its child. Node 3 ('c') has no 'c' ancestors so stays under node 2. Final tree: 0->1->3, 0->2, giving sizes [4,2,1,1].

example_2.py — No Changes Needed

$ Input: parent = [-1,0,0,1,1], s = "abcde"

› Output: [5,3,1,1,1]

💡 Note: All characters are unique, so no node can find an ancestor with the same character. The tree structure remains unchanged, and subtree sizes are calculated normally.

example_3.py — Multiple Relocations

$ Input: parent = [-1,0,1,2,3], s = "aaaaa"

› Output: [5,1,1,1,1]

💡 Note: All nodes have character 'a'. Node 1 stays under root 0. Nodes 2,3,4 all move to be direct children of root 0 (closest 'a' ancestor). Final tree: 0 has children [1,2,3,4].

Constraints

n == parent.length == s.length
1 ≤ n ≤ 10⁵
parent[0] == -1
0 ≤ parent[i] ≤ n - 1 for i ≠ 0
parent represents a valid tree
s consists of only lowercase English letters

Visualization

Tap to expand

DFS Traversal: Visit nodes while maintaining character stacks2. Ancestor Lookup: Check stack top for same character ancestor (O(1))3. Tree Rebuild: Create new parent-child relationships4. Size Calculation: Count subtree sizes in the final structure

Understanding the Visualization

Original Structure

Start with the initial parent-child relationships

Find Matches

Each employee looks upward for a manager with the same specialty

Relocate

Employees move to report directly to their preferred manager

Count Teams

Calculate the size of each manager's team in the new structure

Key Takeaway

🎯 Key Insight: Using character-indexed stacks during DFS enables O(1) ancestor lookup, making the entire algorithm run in optimal O(n) time with just one tree traversal.

Asked in

G Google 23 a Amazon 18 f Meta 15 ⊞ Microsoft 12

The optimal solution uses DFS traversal with character-indexed stacks to efficiently track ancestors. Each character maps to a stack of ancestor nodes on the current path. During DFS, we can find the closest matching ancestor in O(1) time by checking the top of the character's stack. This eliminates the need for repeated upward traversals, achieving O(n) time complexity with a single pass through the tree.

Common Approaches

Approach	Time	Space	Notes
✓ Brute Force (Nested Tree Traversal)	O(n²)	O(n)	For each node, traverse up to find closest matching ancestor, then rebuild tree
DFS with Ancestor Stack (Optimal)	O(n)	O(n)	Use DFS traversal with character-to-node stack mapping for efficient ancestor lookup

Brute Force (Nested Tree Traversal) — Algorithm Steps

Build the original tree structure from parent array
For each node 1 to n-1, traverse up to find closest ancestor with same character
Create new parent relationships based on found ancestors
Build new tree structure with updated parent relationships
Use DFS to calculate subtree sizes for each node

Visualization

Tap to expand

Step-by-Step Walkthrough

Original Tree

Start with the original tree structure

Find Ancestors

Each node searches up for matching character

Rebuild Tree

Create new tree with updated parent relationships

Calculate Sizes

DFS to count subtree sizes

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

int dfs(int node, int** children, int* childCount, int* result) {
    int size = 1;
    for (int i = 0; i < childCount[node]; i++) {
        size += dfs(children[node][i], children, childCount, result);
    }
    result[node] = size;
    return size;
}

int* findSubtreeSizes(int* parent, int parentSize, char* s, int* returnSize) {
    int n = parentSize;
    *returnSize = n;
    
    // Find new parent for each node
    int* newParent = (int*)malloc(n * sizeof(int));
    newParent[0] = -1;
    
    for (int i = 1; i < n; i++) {
        int curr = parent[i];
        while (curr != -1 && s[curr] != s[i]) {
            curr = parent[curr];
        }
        newParent[i] = (curr != -1) ? curr : parent[i];
    }
    
    // Build new tree
    int** newChildren = (int**)malloc(n * sizeof(int*));
    int* childCount = (int*)calloc(n, sizeof(int));
    
    for (int i = 0; i < n; i++) {
        newChildren[i] = (int*)malloc(n * sizeof(int));
    }
    
    for (int i = 1; i < n; i++) {
        newChildren[newParent[i]][childCount[newParent[i]]++] = i;
    }
    
    // Calculate subtree sizes
    int* result = (int*)malloc(n * sizeof(int));
    dfs(0, newChildren, childCount, result);
    
    // Cleanup
    for (int i = 0; i < n; i++) {
        free(newChildren[i]);
    }
    free(newChildren);
    free(childCount);
    free(newParent);
    
    return result;
}

int main() {
    int parent[] = {-1, 0, 1, 2};
    char s[] = "abbc";
    int returnSize;
    int* result = findSubtreeSizes(parent, 4, s, &returnSize);
    for (int i = 0; i < returnSize; i++) {
        printf("%d ", result[i]);
    }
    free(result);
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n²)

For each of n nodes, we may traverse up to n ancestors in worst case

⚠ Quadratic Growth

Space Complexity

O(n)

Space for storing tree structure and recursion stack

⚡ Linearithmic Space

23.4K Views

Medium Frequency

~25 min Avg. Time

847 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Brute Force (Nested Tree Traversal) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler