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

Problem Description

In this problem, we are given a tree consisting of n nodes. Each node has a unique number from 0 to n-1 and a label represented by a lowercase character from the given string labels. The tree has n - 1 edges provided in the edges array, where each element is a pair of nodes indicating there is an edge between those nodes. The goal is to count how many nodes in each node's subtree (including the node itself) have the same label as that particular node and then return the counts in an array.

Intuition

To arrive at the solution for this problem, we should consider a depth-first search (DFS) strategy since we want to deal with subtrees and descendants in a tree structure. The DFS will start from the root node (node 0), explore as far as possible along each branch before backtracking, and perform the necessary counting operations.

Here's the intuition behind the provided solution:

1. We use DFS to traverse the entire tree from the root node, exploring all its subtrees.
2. We maintain a counter, cnt, using a Counter object to track the occurrence of each label as we move down the tree.
3. While performing DFS, for each node, we decrement the count for its label before the DFS goes deeper into its children. This is done to keep track of how many times the label occurs in the current node's subtree.
4. After visiting all of the children nodes, we increment the count back, thus accounting for the occurrence of the node's label in its own subtree.
5. The counter differences before and after exploring the children will give us the number of children with the same label.
6. Create a list g to represent the adjacency list for the graph. For each edge, we add nodes to each other's list to create this undirected graph representation.
7. Create an array ans of size n to store the counts for each node.
8. Call the dfs function with the root node and an invalid parent index -1 to kick-start the DFS.
9. The updated ans array will be returned, which contains the count of nodes in the subtree of every node with the same label as that node.

Learn more about Tree, Depth-First Search and Breadth-First Search patterns.

Solution Approach

The solution uses recursion and graph traversal techniques to address the problem effectively. Here's a breakdown of the solution implementation process with the algorithms, data structures, and patterns employed:

1. Depth-First Search (DFS): A classic DFS recursion is the crux of traversing the tree. The function dfs takes two parameters, i representing the current node, and fa representing the parent of the current node. The recursion helps us reach the leaves, count the nodes, and backtrack to the parent nodes while updating the counts.

2. Counter Object: We use a Counter object from the collections module, named cnt. It keeps track of the count of labels encountered during the DFS traversal. It plays a pivotal role in calculating the number of nodes bearing the same label within a subtree.

3. Adjacency List Representation: A dictionary named g is used to store the adjacency list. Here, we map each node to a list containing its neighboring nodes. This data structure allows us to represent the undirected graph for the tree efficiently.

4. Recursive Logic: In the dfs function:

   • First, we subtract cnt[labels[i]] from ans[i] to record the initial count of the label for node i.
   • We increment cnt[labels[i]] as we've encountered node i's label.
   • We loop over each child j in the adjacency list for node i and perform a DFS call if j is not the parent (to avoid revisiting the parent).
   • After visiting all children, we add back the cnt[labels[i]] count to ans[i], which now contains the correct count for the subtree at i.

5. Result: The ans array, which initially is set to zeros of length n, is updated with the count of nodes that have the same label as node i for each node. This array is returned as the result.

6. Initialization and Start Point: Before starting the DFS traversal, we need to initialize the adjacency list g and the result array ans. We populate g based on the edges provided. The dfs function is called with the root node (0) and -1 as the parent since the root has no parent.

Here is the code block relevant to the above explanation, divided into initialization and recursive DFS parts for clarity:

1# Initialization
2g = defaultdict(list)
3for a, b in edges:
4    g[a].append(b)
5    g[b].append(a)
6cnt = Counter()
7ans = [0] * n
8
9# Recursive DFS Function
10def dfs(i, fa):
11    ans[i] -= cnt[labels[i]]
12    cnt[labels[i]] += 1
13    for j in g[i]:
14        if j != fa:
15            dfs(j, i)
16    ans[i] += cnt[labels[i]]
17
18# Start DFS from root node 0
19dfs(0, -1)

By maintaining a counter that is updated at each node of the DFS, we ensure that when we calculate the count for node i, we have a complete picture of the number of times its label appears within its entire subtree.

Example Walkthrough

Consider a small example where we have a tree with n = 5 nodes. The edges are given by edges = [[0, 1], [0, 2], [1, 3], [1, 4]], and the labels string is "ababa". Using these, we need to count the number of nodes in each node's subtree with the same label as that node.

First, the tree's structure can be visualized as shown below:

1   0(a)
2   / \
3 1(a) 2(b)
4 / \
53(b) 4(a)

According to the labels string, the labels of nodes in the order from 0 to 4 are a, a, b, b, a.

We start with our solution approach:

1. Create the adjacency list g for graph representation:

1g[0] = [1, 2]
2g[1] = [0, 3, 4]
3g[2] = [0]
4g[3] = [1]
5g[4] = [1]

2. Initialize a Counter object, cnt, and an array ans of size n:

1cnt = Counter()
2ans = [0, 0, 0, 0, 0]

3. We start DFS from the root node 0 with -1 as the parent:

During the first call dfs(0, -1):

• ans[0] is decremented by cnt['a'] (which is 0 at this moment), so ans[0] remains 0.

• cnt['a'] is incremented to 1.

• We then call DFS on its children (nodes 1 and 2):

  For node 1:

  • Recursive call dfs(1, 0):
    • ans[1] is decremented by cnt['a'] (which is 1), so ans[1] = -1.
    • cnt['a'] is incremented to 2.
    • dfs(3, 1) is called:
      • ans[3] is decremented by cnt['b'] (0), remains 0.
      • cnt['b'] is incremented to 1.
      • DFS is called on its children (none), so no further action.
      • ans[3] is incremented by cnt['b'] (1), becomes 1.
    • dfs(4, 1) is called:
      • ans[4] is decremented by cnt['a'] (2), becomes -2.
      • cnt['a'] is incremented to 3.
      • DFS is called on its children (none), so no further action.
      • ans[4] is incremented by cnt['a'] (3), becomes 1.
    • After its children are done, ans[1] is incremented by cnt['a'] (3), becomes 2.

  For node 2:

  • Recursive call dfs(2, 0):
    • ans[2] is decremented by cnt['b'] (1), becomes -1.
    • cnt['b'] is incremented to 2.
    • DFS is called on its children (none), so no further action.
    • ans[2] is incremented by cnt['b'] (2), becomes 1.

• After visiting all children of node 0, ans[0] is incremented by cnt['a'] (3), becomes 3.

So, after the DFS completes, we have the ans array updated with the count of the nodes that have the same label as each node in their subtrees:

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

Solution Implementation

Time and Space Complexity

Time Complexity

The time complexity of the code is O(n) where n is the number of nodes in the tree. Each node is visited exactly once due to the depth-first search (DFS). Inside each DFS call, operations are constant time except for the iteration over the children which, across all calls, account for n - 1 edges. Since every edge connects two nodes and each edge is visited exactly twice (once from each node it connects), the time complexity due to edges is O(n - 1) which simplifies to O(n). Thus, the overall time complexity is governed by the number of node visits and edge traversals combined, which is linear in the number of nodes.

Space Complexity

The space complexity is O(n). The dictionary g storing the graph can contain up to 2(n - 1) entries because it is an undirected graph and each edge is stored twice. The cnt Counter and ans array each require O(n) space. The space required for the dfs call stack can also go up to O(n) in the case of a skewed tree where the recursive calls are not returned until the end of the longest branch. Therefore, the total space complexity combines the space for graph representation, counters, answer storage, and the recursion stack, which all scale linearly with the number of nodes in the tree.

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