{x}
blog image

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

You are given a tree (i.e. a connected, undirected graph that has no cycles) consisting of n nodes numbered from 0 to n - 1 and exactly n - 1 edges. The root of the tree is the node 0, and each node of the tree has a label which is a lower-case character given in the string labels (i.e. The node with the number i has the label labels[i]).

The edges array is given on the form edges[i] = [ai, bi], which means there is an edge between nodes ai and bi in the tree.

Return an array of size n where ans[i] is the number of nodes in the subtree of the ith node which have the same label as node i.

A subtree of a tree T is the tree consisting of a node in T and all of its descendant nodes.

 

Example 1:

Input: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], labels = "abaedcd"
Output: [2,1,1,1,1,1,1]
Explanation: Node 0 has label 'a' and its sub-tree has node 2 with label 'a' as well, thus the answer is 2. Notice that any node is part of its sub-tree.
Node 1 has a label 'b'. The sub-tree of node 1 contains nodes 1,4 and 5, as nodes 4 and 5 have different labels than node 1, the answer is just 1 (the node itself).

Example 2:

Input: n = 4, edges = [[0,1],[1,2],[0,3]], labels = "bbbb"
Output: [4,2,1,1]
Explanation: The sub-tree of node 2 contains only node 2, so the answer is 1.
The sub-tree of node 3 contains only node 3, so the answer is 1.
The sub-tree of node 1 contains nodes 1 and 2, both have label 'b', thus the answer is 2.
The sub-tree of node 0 contains nodes 0, 1, 2 and 3, all with label 'b', thus the answer is 4.

Example 3:

Input: n = 5, edges = [[0,1],[0,2],[1,3],[0,4]], labels = "aabab"
Output: [3,2,1,1,1]

 

Constraints:

  • 1 <= n <= 105
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • ai != bi
  • labels.length == n
  • labels is consisting of only of lowercase English letters.

Solution Explanation

This problem asks to count the number of nodes in the subtree of each node that have the same label as the node itself. A Depth-First Search (DFS) approach is ideal for traversing tree structures and counting nodes within subtrees.

Algorithm:

  1. Graph Representation: The input edges array is used to create an adjacency list representation of the tree. This allows efficient traversal of the tree's structure. Each index in the adjacency list represents a node, and the list at that index contains the node's neighbors.

  2. Depth-First Search (DFS): A recursive DFS function (dfs in the code) is used to explore the tree. The function takes the current node (i), its parent (fa), and a counter array (cnt) that tracks the count of each character label encountered within the current subtree.

  3. Counting Labels: For each node visited in the DFS, the following steps are performed:

    • Subtract the current count of the node's label (cnt[labels[i]]) from ans[i]. This is a crucial step for accurate counting. It ensures that when we encounter a node, we subtract the count of that label from its subtree. Then, we increase the count for the character of the current node's label in cnt.
    • Recursively call dfs for all children of the current node (excluding the parent to avoid cycles). This explores the entire subtree rooted at the current node.
    • After exploring the children, add the count of the node's label (cnt[labels[i]]) back to ans[i]. This completes the count of the nodes with the same label as the current node in its subtree.

  4. Result: The ans array, initialized with zeros, is used to store the result. ans[i] will contain the count of nodes with the same label as node i within its subtree after the DFS completes.

Time Complexity Analysis:

The DFS function visits each node and edge exactly once. The time complexity is directly proportional to the number of nodes and edges in the tree. Since the tree has n nodes and n-1 edges, the time complexity is O(n).

Space Complexity Analysis:

The space complexity is dominated by the recursion stack in the DFS function. In the worst-case scenario (a skewed tree), the recursion depth can be O(n). Additionally, the adjacency list and the cnt and ans arrays require O(n) space. Therefore, the overall space complexity is O(n).

Code Explanation (Python):

from collections import defaultdict, Counter
 
class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
        # Create adjacency list
        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
 
        # Initialize counter and result arrays
        cnt = Counter()
        ans = [0] * n
 
        def dfs(i, fa):  # DFS function
            k = labels[i]  # Get current node's label
            ans[i] -= cnt[k] # Subtract the current count of the node's label from its subtree count
            cnt[k] += 1     # Increase the count of the current node's label
 
            for j in g[i]:
                if j != fa:
                    dfs(j, i)   # Recursive call for children
 
            ans[i] += cnt[k]   # Add the updated count of the node's label back into its subtree count
 
        dfs(0, -1)  # Start DFS from root (node 0)
        return ans

The other language implementations follow a similar structure and logic, adapting syntax and data structures specific to each language. The core algorithm remains the same efficient DFS-based approach.