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

Number of Nodes in the Sub-Tree With the Same Label

Number: 1643

Difficulty: Medium

Paid? No

Companies: Samsung, Uber

Problem Description

Given a tree of n nodes (indexed from 0 to n - 1) where node 0 is the root, each node has a label represented by a lowercase character in the string labels. The tree is represented using an array of edges, and there is exactly n - 1 edges. The task is to return an array ans of length n where ans[i] is the number of nodes in the subtree of node i (including itself) that share the same label as node i.

Key Insights

  • The structure is guaranteed to be a tree (a connected, acyclic graph), allowing use of DFS without worrying about cycles.
  • For every node, we need to count the occurrences of its own label in its subtree.
  • Use DFS recursion with an adjacency list to traverse the tree and aggregate label frequencies from child nodes.
  • Maintain a frequency count (e.g., using an array of size 26 for each lowercase letter) for each subtree and update the parent's counts accordingly.

Space and Time Complexity

Time Complexity: O(n) – each node and each edge is processed once. Space Complexity: O(n) – for the recursion stack and the adjacency list to represent the tree.

Solution

We solve the problem using Depth-First Search (DFS). First, construct an adjacency list from the given edges. Start DFS from the root node (0), and for each node initialize a frequency array for the 26 lowercase characters. Recursively, visit each child (skipping the parent to prevent revisiting) and accumulate the frequency counts returned from each subtree. After processing all children, update the frequency of the current node’s label, record this count as ans[node], and finally, return the frequency array to be merged with the parent's counts.

Code Solutions

# Build the adjacency list for the graph
from collections import defaultdict, Counter

def countSubTrees(n, edges, labels):
    # Create an adjacency list to represent the tree
    graph = defaultdict(list)
    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)

    # Initialize the result list to store the answer for each node
    result = [0] * n

    # Define the DFS function
    def dfs(current, parent):
        # Create a frequency counter array (using list for 26 letters)
        count = [0] * 26
        # Process all children except the parent to avoid cycles
        for neighbor in graph[current]:
            if neighbor == parent:
                continue
            # Recursively collect counts from child subtrees
            childCount = dfs(neighbor, current)
            # Merge the frequency counts from the child with the current count
            for i in range(26):
                count[i] += childCount[i]
                
        # Update count for the current node's label
        labelIndex = ord(labels[current]) - ord('a')
        count[labelIndex] += 1
        # Save the count for the current node in the result array
        result[current] = count[labelIndex]
        # Return the frequency count array upward for merging
        return count

    # Start DFS from the root node (0) with an invalid parent (-1)
    dfs(0, -1)
    return result

# Example usage:
#print(countSubTrees(7, [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], "abaedcd"))
← Back to All Questions