Find Subtree Sizes After Changes

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

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

✓ Brute Force (Nested Tree Traversal)

⏱️ Time: O(n²) Space: O(n)

For each node from 1 to n-1, we traverse up the tree to find the closest ancestor with the same character. Then we build the new tree structure and calculate subtree sizes using DFS.

DFS with Ancestor Stack (Optimal)

⏱️ Time: O(n) Space: O(n)

The optimal approach uses a single DFS traversal while maintaining a stack for each character that tracks the current path of ancestors. This allows O(1) lookup of the closest ancestor with any given character.

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>

static int** newChildren;
static int* childCount;
static int* result;

int dfs(int node) {
    int size = 1;
    for (int i = 0; i < childCount[node]; i++) {
        size += dfs(newChildren[node][i]);
    }
    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
    newChildren = (int**)malloc(n * sizeof(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
    result = (int*)malloc(n * sizeof(int));
    dfs(0);
    
    // Cleanup
    for (int i = 0; i < n; i++) {
        free(newChildren[i]);
    }
    free(newChildren);
    free(childCount);
    free(newParent);
    
    return result;
}

int main() {
    char line[1000];
    char s[1000];
    
    fgets(line, sizeof(line), stdin);
    fgets(s, sizeof(s), stdin);
    
    // Remove newline
    line[strcspn(line, "\n")] = 0;
    s[strcspn(s, "\n")] = 0;
    
    // Parse parent array
    int parent[1000];
    int n = 0;
    char* token = strtok(line, ",");
    while (token != NULL) {
        parent[n++] = atoi(token);
        token = strtok(NULL, ",");
    }
    
    int returnSize;
    int* result = findSubtreeSizes(parent, n, s, &returnSize);
    
    for (int i = 0; i < returnSize; i++) {
        if (i > 0) printf(",");
        printf("%d", result[i]);
    }
    printf("\n");
    
    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