Facebook Pixel

3331. Find Subtree Sizes After Changes

MediumTreeDepth-First SearchArrayHash TableString
Leetcode Link

Problem Description

You are given a tree rooted at node 0, consisting of n nodes numbered from 0 to n - 1. This tree is represented by an array parent of size n, where parent[i] is the parent of node i. Node 0 is the root, so parent[0] is -1. Additionally, you are provided a string s of length n, where s[i] is the character assigned to node i.

The task is to make changes to the tree one time and simultaneously for all nodes from 1 to n - 1 as follows:

  • Find the closest ancestor y of node x such that s[x] equals s[y].
  • If no such ancestor exists, do nothing.
  • Otherwise, remove the edge between x and its current parent and make y the new parent of x.

Your goal is to return an array answer of size n where answer[i] is the size of the subtree rooted at node i in the final tree after all transformations have been completed.

Intuition

The key idea is to keep track of the nearest ancestor of each node that has the same character. This can be efficiently managed using Depth First Search (DFS) by maintaining a dictionary where keys are characters and values are lists of nodes visited in that DFS traversal.

The DFS starts at the root node (node 0). For each node, we:

  1. Record the node in a dictionary based on the character at that node (s[i]).
  2. Recursively perform DFS on its children to calculate subtree sizes.
  3. After processing the children of a node, check if the node has a same-character ancestor stored in the dictionary. If yes, this ancestor becomes the new parent of the current node.
  4. Calculate the subtree size by accumulating sizes from children.
  5. Once finished with the current node, remove it from the dictionary to maintain only the most recent positions of each character.

This approach ensures that each node's final subtree size is computed after restructuring the tree based on same-character ancestry. The final answer is built up by accumulating computed subtree sizes in an answer array.

Learn more about Tree and Depth-First Search patterns.

Solution Approach

The solution employs a Depth First Search (DFS) approach to navigate through the tree. Here's a step-by-step breakdown of the implementation:

  1. Data Structure Initialization:

    • Construct a graph g from the parent array, representing the tree structure. Each index in list g holds the children of corresponding node.
    • Use a default dictionary d to record nodes by character as they are encountered in the DFS.
    • Initialize the answer array ans to store the size of the subtree for each node.

  2. DFS Function:

    • Define a recursive function dfs(i, fa) where i is the current node and fa is its parent.
    • For each node i:
      • Start with ans[i] = 1 because the node itself counts toward the subtree size.
      • Append node i to d[s[i]], storing the node in the list associated with its character.
      • Recursively call dfs on each child of node i.
      • After processing children, attempt to find the closest ancestor with the same character as node i. This is done by looking at the second-to-last entry in d[s[i]] because the last entry is the node itself.
      • If such an ancestor exists (k != -1), increment the ancestor's subtree size by the size of the current node's subtree.

  3. Post-DFS Cleanup:

    • After processing a node and its subtree, remove the node from d[s[i]] to maintain the correct state of the dictionary.

  4. Return Result:

    • Once the DFS completes from the root node, the array ans holds the size of the subtree rooted at each node after restructuring. Return this as the final result.

This approach uses DFS for traversal and leverages a dictionary to track ancestorship by character, ensuring each node correctly aligns itself to the nearest ancestor with a matching character, thereby forming the desired tree structure.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's consider a simple example to illustrate the solution approach.

Example:

  • Input:
    • parent = [-1, 0, 0, 1, 1, 2, 2]
    • s = "abacada"

Here, parent represents a tree rooted at node 0 with n = 7 nodes, and s assigns each node a character. The tree structure is as follows:

  • 0 is the root.
  • 1 and 2 are children of 0.
  • 3 and 4 are children of 1.
  • 5 and 6 are children of 2.

Step-by-Step Solution:

  1. Data Structure Initialization:

    • Construct a graph g from the parent array: 
      g = [[1, 2], [3, 4], [5, 6], [], [], [], []]
    • Initialize a default dictionary d to track nodes by character.
    • Prepare an answer array ans initialized with zeros: 
      ans = [0, 0, 0, 0, 0, 0, 0]

  2. DFS Traversal:

    • Begin DFS at root node 0.

    • For node 0:

      • Set ans[0] = 1 (includes itself).
      • Record node 0 in d under key 'a': 
        d = {'a': [0]}
      • Visit children 1 and 2.

    • For node 1:

      • Set ans[1] = 1.
      • Record node 1 in d under key 'b': 
        d = {'a': [0], 'b': [1]}
      • Visit its children 3 and 4.

    • For node 3:

      • Set ans[3] = 1.
      • Add node 3 under key 'a': 
        d = {'a': [0, 3], 'b': [1]}
      • Find closest 'a' ancestor for 3 (node 0) and adjust subtree: new parent 0, update ans[0].
      • Remove node 3 from d: 
        d = {'a': [0], 'b': [1]}

    • For node 4:

      • Set ans[4] = 1.
      • Add node 4 under key 'c': 
        d = {'a': [0], 'b': [1], 'c': [4]}
      • No same-character ancestor for 4, so no changes.
      • Remove node 4: 
        d = {'a': [0], 'b': [1], 'c': []}

    • Update ans[1] after children:

      ans = [1, 2, 0, 1, 1, 0, 0]

    • For node 2:

      • Set ans[2] = 1.
      • Add node 2 under 'a': 
        d = {'a': [0, 2], 'b': [1]}
      • Visit children 5 and 6.

    • For node 5:

      • Set ans[5] = 1.
      • Add node 5 under 'd': 
        d = {'a': [0, 2], 'b': [1], 'd': [5]}
      • No same-character ancestor for 5, so no changes.
      • Remove node 5: 
        d = {'a': [0, 2], 'b': [1], 'd': []}

    • For node 6:

      • Set ans[6] = 1.
      • Add node 6 under 'a': 
        d = {'a': [0, 2, 6], 'b': [1]}
      • Find closest 'a' ancestor (node 2) and restructure tree.
      • Clean up: 
        d = {'a': [0, 2], 'b': [1]}

    • Update ans[2] after children:

      ans = [1, 2, 2, 1, 1, 1, 1]

  3. Return Result:

    • The resulting ans array reflects the size of each subtree: 
      answer = [7, 3, 3, 1, 1, 1, 1]
    • Return this as the final answer.

Solution Implementation

1from typing import List
2from collections import defaultdict
3
4class Solution:
5    def findSubtreeSizes(self, parent: List[int], s: str) -> List[int]:
6        def dfs(node: int, parent_node: int):
7            # Initialize the subtree size for the current node to 1
8            subtree_sizes[node] = 1
9            # Append the current node index to the list of nodes for character s[node]
10            character_dict[s[node]].append(node)
11            # Recursively traverse the children of the current node
12            for child in graph[node]:
13                dfs(child, node)
14            # Determine the last visited ancestor with the same character
15            ancestor = parent_node
16            if len(character_dict[s[node]]) > 1:
17                ancestor = character_dict[s[node]][-2]
18            # If a valid ancestor is found, update its subtree size
19            if ancestor != -1:
20                subtree_sizes[ancestor] += subtree_sizes[node]
21            # Remove the current node from the character's list (backtracking)
22            character_dict[s[node]].pop()
23
24        n = len(s)
25        # Create an adjacency list representation of the tree
26        graph = [[] for _ in range(n)]
27        for i in range(1, n):
28            graph[parent[i]].append(i)
29        # Dictionary to keep track of nodes with the same character
30        character_dict = defaultdict(list)
31        # List to store the size of the subtree rooted at each node
32        subtree_sizes = [0] * n
33        # Start DFS traversal from the root node (0) with no parent (-1)
34        dfs(0, -1)
35        return subtree_sizes
36
1import java.util.List;
2import java.util.ArrayList;
3import java.util.Arrays;
4
5class Solution {
6    private List<Integer>[] adjacencyList;  // Represents graph edges as adjacency list
7    private List<Integer>[] characterIndices;  // Stores indices of characters in the string
8    private char[] charArray;  // Character array of the input string
9    private int[] subtreeSizes;  // Array to hold resulting subtree sizes
10
11    // Method to find subtree sizes for each node based on string properties
12    public int[] findSubtreeSizes(int[] parent, String s) {
13        int n = s.length();  // Total number of nodes
14
15        adjacencyList = new List[n];  // Initialize the adjacency list
16        characterIndices = new List[26];  // Initialize arrays to store character indices ('a' to 'z')
17
18        charArray = s.toCharArray();  // Convert string to character array
19
20        // Initialize each element of adjacency list and characterIndices list
21        Arrays.setAll(adjacencyList, k -> new ArrayList<>());
22        Arrays.setAll(characterIndices, k -> new ArrayList<>());
23
24        // Build the adjacency list from parent-child relationships
25        for (int i = 1; i < n; ++i) {
26            adjacencyList[parent[i]].add(i);  // Add child to parent's list
27        }
28
29        subtreeSizes = new int[n];  // Initialize subtree sizes array
30
31        dfs(0, -1);  // Start DFS from root node
32
33        return subtreeSizes;  // Return computed subtree sizes
34    }
35
36    // Depth First Search recursive method to compute subtree sizes
37    private void dfs(int currentNode, int parentNode) {
38        subtreeSizes[currentNode] = 1;  // Start with size 1 for the node itself
39        int charIndex = charArray[currentNode] - 'a';  // Find the character index (0 for 'a', 1 for 'b', etc.)
40      
41        characterIndices[charIndex].add(currentNode);  // Add current node to the list of indices for its character
42
43        // Recursively traverse each child node
44        for (int childNode : adjacencyList[currentNode]) {
45            dfs(childNode, currentNode);  // DFS on child
46        }
47
48        // Fetch the 'parent node' in the context of the current character path
49        int previousNode = characterIndices[charIndex].size() > 1 ? 
50                characterIndices[charIndex].get(characterIndices[charIndex].size() - 2) : parentNode;
51      
52        // If valid previous node exists, add current subtree size to it
53        if (previousNode >= 0) {
54            subtreeSizes[previousNode] += subtreeSizes[currentNode];
55        }
56
57        // Remove current node from the list of character indices
58        characterIndices[charIndex].remove(characterIndices[charIndex].size() - 1);
59    }
60}
61
1#include <vector>
2#include <string>
3using namespace std;
4
5class Solution {
6public:
7    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
8        int n = s.size();
9      
10        // Graph to store children of each node
11        vector<vector<int>> graph(n);
12      
13        // An array to keep track of nodes for each character in the string
14        vector<vector<int>> charNodes(26);
15      
16        // Build the graph from the parent list
17        for (int i = 1; i < n; ++i) {
18            graph[parent[i]].push_back(i);
19        }
20      
21        // Resultant vector to store subtree sizes
22        vector<int> result(n);
23      
24        // Recursive lambda function for Depth First Search
25        function<void(int, int)> dfs = [&](int node, int parent) {
26            // Initialize current node size as 1 (counting itself)
27            result[node] = 1;
28          
29            // Get the character index for the current node
30            int charIndex = s[node] - 'a';
31          
32            // Add the current node to the corresponding character's list
33            charNodes[charIndex].push_back(node);
34          
35            // Recursively visit each child
36            for (int child : graph[node]) {
37                dfs(child, node);
38            }
39          
40            // Find the parent node for this character or default to the given parent
41            int prevNode = charNodes[charIndex].size() > 1 ? charNodes[charIndex][charNodes[charIndex].size() - 2] : parent;
42          
43            // Add the current subtree size to the previous node in the character's list
44            if (prevNode >= 0) {
45                result[prevNode] += result[node];
46            }
47          
48            // Remove the current node from the character's list
49            charNodes[charIndex].pop_back();
50        };
51      
52        // Start DFS from the root node with no parent (-1)
53        dfs(0, -1);
54      
55        // Return the result containing sizes of subtrees with the same character
56        return result;
57    }
58};
59
1function findSubtreeSizes(parent: number[], s: string): number[] {
2    const n = parent.length;  // Number of nodes in the tree
3    const g: number[][] = Array.from({ length: n }, () => []); // Adjacency list for the tree
4    const d: number[][] = Array.from({ length: 26 }, () => []); // Array to keep track of character occurrences
5
6    // Build the tree using the parent array
7    for (let i = 1; i < n; ++i) {
8        g[parent[i]].push(i);
9    }
10
11    const ans: number[] = Array(n).fill(1); // Initialize subtree sizes, each node itself is a subtree
12
13    // Depth First Search to compute subtree sizes
14    const dfs = (current: number, predecessor: number): void => {
15        const charIndex = s.charCodeAt(current) - 97; // Convert character to index
16        d[charIndex].push(current); // Add current node to the character index list
17
18        // Explore children nodes
19        for (const child of g[current]) {
20            dfs(child, current);
21        }
22
23        // Find previous node with the same character or use predecessor
24        const prev = d[charIndex].length > 1 ? d[charIndex][d[charIndex].length - 2] : predecessor;
25      
26        if (prev >= 0) {
27            ans[prev] += ans[current]; // Update the size of the subtree
28        }
29
30        d[charIndex].pop(); // Remove current node from the character index list
31    };
32
33    dfs(0, -1); // Start DFS from the root node
34    return ans; // Return calculated subtree sizes
35}
36

Time and Space Complexity

The time complexity of the code is O(n), where n is the number of nodes in the tree. This is because the depth-first search (DFS) traverses each node once, and each operation within the DFS (such as updating lists or dictionaries) takes constant time.

The space complexity of the code is also O(n). This results from storing the tree structure in the adjacency list g, the dictionary d to maintain characters at each depth, and the ans list to store subtree sizes. Additionally, the recursion stack used during the DFS can go as deep as n in the worst case.

Learn more about how to find time and space complexity quickly.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

Which two pointer techniques do you use to check if a string is a palindrome?

Recommended Readings

Everything About Trees

Everything About Trees A tree is a type of graph data structure composed of nodes and edges Its main properties are It is acyclic doesn't contain any cycles There exists a path from the root to any node Has N 1 edges where N is the number of nodes in the tree and

DFS Fundamentals

https algomonster s3 us east 2 amazonaws com cover_photos dfs svg Depth First Search Prereqs Recursion Review problems recursion_intro Trees problems tree_intro With a solid understanding of recursion under our belts we are now ready to tackle one of the most useful techniques in coding interviews Depth First Search DFS

Coding Interview Patterns: The Smart, Data-Driven Way to Prepare for Tech Interviews

Coding Interview Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way

Want a Structured Path to Master System Design Too? Don’t Miss This!

Load More