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
✓ 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)	O(n²)	O(n)	For each node, traverse up to find closest matching ancestor, then rebuild tree

DFS with Ancestor Stack (Optimal) — Algorithm Steps

Initialize a hash map where each character maps to a stack of ancestor nodes
Start DFS from root, maintaining ancestor stacks during traversal
For each node, check if there's an ancestor with the same character in the stack
If found, reassign the node to that ancestor; otherwise keep original parent
Update ancestor stacks as we traverse (push on entry, pop on exit)
Calculate subtree sizes during the same DFS traversal

Visualization

Tap to expand

Step-by-Step Walkthrough

Initialize

Start DFS with empty character stacks

Push & Check

For each node, check stack for same character ancestor

Process

Add node to its character stack, process children

Backtrack

Pop node from stack when returning from recursion

Code -

solution.c — C

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

#define MAX_NODES 100000
#define MAX_CHARS 26

typedef struct {
    int* stack;
    int top;
    int capacity;
} CharStack;

CharStack charStacks[MAX_CHARS];
int** children;
int* childCount;
int* newParent;
int* result;
int* parent;
char* s;

void initStack(CharStack* stack) {
    stack->stack = (int*)malloc(MAX_NODES * sizeof(int));
    stack->top = -1;
    stack->capacity = MAX_NODES;
}

void push(CharStack* stack, int val) {
    if (stack->top < stack->capacity - 1) {
        stack->stack[++stack->top] = val;
    }
}

void pop(CharStack* stack) {
    if (stack->top >= 0) {
        stack->top--;
    }
}

int top(CharStack* stack) {
    return stack->top >= 0 ? stack->stack[stack->top] : -1;
}

int isEmpty(CharStack* stack) {
    return stack->top < 0;
}

int dfs(int node) {
    // Find closest ancestor with same character
    if (node != 0) {
        int charIdx = s[node] - 'a';
        if (!isEmpty(&charStacks[charIdx])) {
            newParent[node] = top(&charStacks[charIdx]);
        } else {
            newParent[node] = parent[node];
        }
    }
    
    // Add current node to its character stack
    int charIdx = s[node] - 'a';
    push(&charStacks[charIdx], node);
    
    // Process children
    int subtreeSize = 1;
    for (int i = 0; i < childCount[node]; i++) {
        subtreeSize += dfs(children[node][i]);
    }
    
    // Remove current node from character stack
    pop(&charStacks[charIdx]);
    
    result[node] = subtreeSize;
    return subtreeSize;
}

int calcSizes(int node, int** newChildren, int* newChildCount) {
    int size = 1;
    for (int i = 0; i < newChildCount[node]; i++) {
        size += calcSizes(newChildren[node][i], newChildren, newChildCount);
    }
    result[node] = size;
    return size;
}

int* findSubtreeSizes(int* parentArr, int parentSize, char* str, int* returnSize) {
    int n = parentSize;
    *returnSize = n;
    
    parent = parentArr;
    s = str;
    
    // Initialize character stacks
    for (int i = 0; i < MAX_CHARS; i++) {
        initStack(&charStacks[i]);
    }
    
    // Build original tree
    children = (int**)malloc(n * sizeof(int*));
    childCount = (int*)calloc(n, sizeof(int));
    newParent = (int*)malloc(n * sizeof(int));
    result = (int*)malloc(n * sizeof(int));
    
    for (int i = 0; i < n; i++) {
        children[i] = (int*)malloc(n * sizeof(int));
        newParent[i] = -1;
    }
    
    for (int i = 1; i < n; i++) {
        children[parent[i]][childCount[parent[i]]++] = i;
    }
    
    dfs(0);
    
    // Rebuild tree with new parent relationships
    int** newChildren = (int**)malloc(n * sizeof(int*));
    int* newChildCount = (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++) {
        if (newParent[i] != -1) {
            newChildren[newParent[i]][newChildCount[newParent[i]]++] = i;
        }
    }
    
    calcSizes(0, newChildren, newChildCount);
    
    // Cleanup
    for (int i = 0; i < MAX_CHARS; i++) {
        free(charStacks[i].stack);
    }
    for (int i = 0; i < n; i++) {
        free(children[i]);
        free(newChildren[i]);
    }
    free(children);
    free(newChildren);
    free(childCount);
    free(newChildCount);
    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)

Single DFS traversal visiting each node exactly once

✓ Linear Growth

Space Complexity

O(n)

Space for ancestor stacks and recursion depth, both bounded by tree height and character set

⚡ 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

DFS with Ancestor Stack (Optimal) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler