Program to find number of nodes in the sub-tree with the same label using Python

In this problem, we have a rooted tree with n nodes numbered from 0 to n-1, where each node has a label (lowercase English letter). We need to find the count of nodes in each node's subtree that have the same label as that node.

Problem Understanding

Given a tree structure and node labels, we need to return an array where each element represents the count of nodes with the same label in that node's subtree (including the node itself).

For the tree above with labels "ccaca", the result is [3, 2, 1, 1, 1] because ?

• Node 0 (label 'c'): has 3 nodes with label 'c' in its subtree (0, 1, 3)
• Node 1 (label 'c'): has 2 nodes with label 'c' in its subtree (1, 3)
• Nodes 2, 3, 4: each has 1 node with their respective label in their subtree

Algorithm

We'll use a DFS approach with the following steps ?

1. Build an adjacency list from the edge list
2. Use DFS to traverse the tree and count labels in each subtree
3. For each node, count occurrences of its label in its subtree

Implementation

from collections import defaultdict, Counter

def count_sub_tree_with_same_label(n, edges, labels):
    # Build adjacency list
    graph = defaultdict(set)
    for parent, child in edges:
        graph[parent].add(child)
        graph[child].add(parent)
    
    # Result array to store counts
    result = [0] * n
    
    def dfs(node):
        # Counter to track label frequencies in current subtree
        label_count = Counter()
        
        # Process all children
        for neighbor in graph[node]:
            # Remove edge to prevent revisiting parent
            graph[neighbor].remove(node)
            # Get label counts from child subtree
            child_counts = dfs(neighbor)
            # Merge child counts with current counts
            label_count.update(child_counts)
        
        # Add current node's label
        label_count[labels[node]] += 1
        
        # Store count for current node's label
        result[node] = label_count[labels[node]]
        
        return label_count
    
    # Start DFS from root node (0)
    dfs(0)
    return result

# Test the function
n = 5
edges = [[0,1], [0,2], [1,3], [0,4]]
labels = "ccaca"
print(count_sub_tree_with_same_label(n, edges, labels))

[3, 2, 1, 1, 1]

How It Works

The algorithm performs a post-order DFS traversal ?

1. Graph Construction: Build an undirected graph from the edge list
2. DFS Traversal: Start from root (node 0) and visit all children
3. Label Counting: For each node, count labels in its subtree using a Counter
4. Edge Removal: Remove visited edges to prevent cycles during traversal
5. Result Storage: Store the count of the node's label in its subtree

Time and Space Complexity

• Time Complexity: O(n) - we visit each node exactly once
• Space Complexity: O(n) - for the adjacency list, recursion stack, and result array

Conclusion

This solution uses DFS with label counting to efficiently find the number of nodes with the same label in each subtree. The key insight is using post-order traversal to aggregate counts from children before processing the current node.