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

Problem Explanation

The problem is about finding the number of nodes in subtrees of a given tree that have the same label as the root of that subtree. The tree is defined here as a connected, undirected graph that has no cycles and the tree consists of n nodes numbered 0 to n - 1 with exactly n - 1 edges. A label for each node is given in the form of a string of lowercase characters. The goal is to return an array of size n where the i-th element represents the number of nodes in the subtree rooted at node i that have the same label as node i.

Example

Let's walk through Example 1:

We have a tree with 7 nodes (n = 7), the edges are [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], and the labels are "abaedcd". Now let's see the subtree of each node:

  • Node 0: Subtree includes nodes 0,1,2,4,5,3,6. Nodes 0 and 2 have the same label 'a'. So, the count is 2.
  • Node 1: Subtree include nodes 1,4,5. Only node 1 has the label 'b'. So, the count is 1.
  • Node 2: Subtree include nodes 2,3,6. Only node 2 has the label 'c'. So, the count is 1.
  • Node 3,4,5,6: Each of these nodes has no children. So, they each count for 1.

The final output is [2,1,1,1,1,1,1].

Solution Approach

To solve this problem, we need to traverse the tree and count the labels for each sub-tree. A Depth First Search (DFS) approach is suitable for traversing the tree from root to each leaf.

During the DFS, for each node, we create an array to store the counts for every label (26 lowercase characters) in its sub-tree including itself. Then, for every child of the current node, we update the counts array by adding the count of each label in its sub-tree. In the end, we increase the count for the label of the current node by 1, as the node itself is part of its subtree.

Let's walk through the solution using above example.

  1. We start DFS from the root node, which is node 0. Initially empty count array is used to store the label count.

  2. Then, we go to its first child which is node 1 (node 1 is child of node 0 as there is an edge between node 0 and node 1). Here, for node 1 we repeat the same process.

  3. In the subtree rooted at node 1, we visit its children nodes 4 and 5. Since these nodes do not have further children, we return their counts (with the number of occurrences of their label incremented by 1). The total count for node 1 therefore becomes 1 as only node 1 has label 'b'.

  4. We continue same for all nodes and finally return the count array.

This process can be performed efficiently using an adjacency list representation for the graph and then performing DFS. The solution iterates over each node once, giving it a time complexity of O(n).

Python Solution

1
2python
3from collections import defaultdict
4
5class Solution:
6    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
7        graph = defaultdict(list)
8        for u, v in edges:
9            graph[u].append(v)
10            graph[v].append(u)
11
12        res = [0] * n
13        # DFS function to calculate and update counts
14        def dfs(node, parent):
15            count = [0] * 26
16            # ASCII value of current label minus 'a'
17            index = ord(labels[node]) - ord('a')
18            count[index] = 1 
19            # Add counts of child labels to current node's count
20            for child in graph[node]:
21                if child == parent: 
22                    continue
23                child_count = dfs(child, node)
24                for i in range(26):
25                    count[i] += child_count[i]
26            res[node] = count[index] 
27            return count
28		
29        dfs(0, -1)
30        return res

Java Solution

1
2java
3class Solution {
4     public int[] countSubTrees(int n, List<List<Integer>> edges, String labels) {
5        List<Integer>[] graph = new List[n];
6        for (int i = 0; i < n; ++i)
7            graph[i] = new ArrayList();
8        for (List<Integer> edge : edges) {
9            int u = edge.get(0), v = edge.get(1);
10            graph[u].add(v);
11            graph[v].add(u);
12        }
13        int[] ans = new int[n];
14        boolean[] seen = new boolean[n];
15        dfs(graph, 0, labels, ans, seen);
16        return ans;
17    }
18
19    private int[] dfs(List<Integer>[] graph, int u, String labels, int[] ans, boolean[] seen) {
20        int[] count = new int[26];
21        if (!seen[u]) {
22            seen[u] = true;
23            for (int v: graph[u]) {
24                int[] res = dfs(graph, v, labels, ans, seen);
25                for (int i = 0; i < 26; ++i)
26                    count[i] += res[i];
27            }
28            count[labels.charAt(u)-'a']++;
29            ans[u] = count[labels.charAt(u)-'a'];
30        }
31        return count;
32    }
33}

Javascript Solution

1
2javascript
3var countSubTrees = function(n, edges, labels) {
4    let g = [...Array(n)].map(d => Array()), seen = new Set(), res = Array(n);
5    for (let [u, v] of edges)
6        g[u].push(v), g[v].push(u);
7    let dfs = function(v) {
8        let cnt = new Array(26).fill(0);
9        seen.add(v);
10        for (let u of g[v])
11            if (!seen.has(u))
12                let cur = dfs(u);
13                for (let i = 0; i < 26; ++i)
14                    cnt[i] += cur[i];
15        cnt[labels.codePointAt(v) - 97]++;
16        res[v] = cnt[labels.codePointAt(v) - 97];
17        return cnt;
18    };
19    dfs(0);
20    return res;
21};
22

C++ Solution

1
2cpp
3class Solution {
4public:
5    vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
6        vector<vector<int>> tree(n);
7        for(const vector<int>& edge : edges){
8            int u = edge[0], v = edge[1];
9            tree[u].push_back(v);
10            tree[v].push_back(u);
11        }
12        vector<int> ans(n);
13        dfs(0,-1,tree,labels, ans);
14        return ans;
15    }
16
17private:
18    vector<int> dfs(int u, int parent, const vector<vector<int>>& tree, const string& labels, vector<int>& ans){
19        vector<int> count(26, 0);
20        for(int v : tree[u]){
21            if(v == parent) continue;
22            vector<int> tempCount = dfs(v, u, tree, labels, ans);
23            for(int i = 0; i<26; ++i)
24                count[i]+= tempCount[i];
25        }
26        ans[u] = ++count[labels[u]-'a'];
27        return count;
28    }
29};

C# Solution

1
2csharp
3public class Solution {
4    public int[] CountSubTrees(int n, int[][] edges, string labels) {
5        List<int>[] graph = new List<int>[n];
6        for (int i = 0; i < n; i++) 
7            graph[i] = new List<int>();
8        foreach (var edge in edges) {
9            graph[edge[0]].Add(edge[1]);
10            graph[edge[1]].Add(edge[0]);
11        }
12
13        int[] result = new int[n];
14        bool[] seen = new bool[n];
15        DFS(0, graph, seen, labels, result);
16        return result;
17    }
18    
19    private int[] DFS(int node, List<int>[] graph, bool[] seen, string labels, int[] result) {
20        int[] count = new int[26];
21        if (!seen[node]) {
22            seen[node] = true;
23            foreach (var i in graph[node]) {
24                var res = DFS(i, graph, seen, labels, result);
25                for (int j = 0; j < 26; j++)
26                    count[j] += res[j];
27            }
28            count[labels[node]-'a']++;
29            result[node] = count[labels[node]-'a'];
30        }        
31        return count;
32    }
33}
34}

Phew! That was a lot of code examples in various programming languages. As you can see, each solution essentially consists of the same steps and logic. The differences are due to syntax and language-specific features.

Let's recap the solution:

  1. Create a Graph: We need to convert the given edge list to an adjacency list to use as a tree.

    In Python, we use defaultdict to create a dictionary where each key corresponds to a node, and the value is a list of its children nodes.

    In Java and C#, we use an array of ArrayList.

    In JavaScript, we use an array of arrays.

    In C++, we use a vector of vectors.

  2. Initialize the Result Array: This array will hold the count of labels for each node.

  3. Define the DFS function: This recursive function is where the actual computation is done. It takes a node and its parent node. For each child of the current node, it updates the counts array by adding the counts of labels from the child node.

    Then, it increments the count of the current node's label (since the node itself is part of its subtree). The updated counts array is stored in the result array for the current node.

  4. Call DFS function: We start the DFS traversal from the root node which is node 0.

  5. Return the result: The final result, an array of counts for each node, is then returned.

Remember, the depth-first search (DFS) traversal is key to the solution because it enables visiting each node's children before backtracking, which is ideal for this problem since the count of a node depends on the counts of its child nodes.

Thus, the code in all those different languages employs DFS to traverse the nodes and tracks the counts of labels using an array or a list. With slight variations, the DFS function and the counts array setup are done, making this problem a good study in comparing coding constructs and approaches across languages.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫