Problem

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.

Examples

Example 1:

1
2
3
4
5
6

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:

1
2
3
4
5
6
7
8

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:

1
2
3
4

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

Solution

Method 1 – DFS with Frequency Counting

Intuition

To count the number of nodes in each subtree with the same label, we use DFS to traverse the tree. For each node, we maintain a frequency array of labels in its subtree and update the answer for that node using its label’s frequency.

Approach

  1. Build the tree as an adjacency list from the edges.
  2. Define a DFS function that returns a frequency array for each subtree.
  3. For each node, recursively merge the frequency arrays of its children.
  4. Increment the frequency for the current node’s label and set the answer for that node.
  5. Return the answer array after DFS traversal.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

class Solution {
    public int[] countSubTrees(int n, int[][] edges, String labels) {
        List<Integer>[] tree = new List[n];
        for (int i = 0; i < n; i++) tree[i] = new ArrayList<>();
        for (int[] e : edges) {
            tree[e[0]].add(e[1]);
            tree[e[1]].add(e[0]);
        }
        int[] ans = new int[n];
        dfs(0, -1, tree, labels, ans);
        return ans;
    }
    private int[] dfs(int u, int p, List<Integer>[] tree, String labels, int[] ans) {
        int[] freq = new int[26];
        freq[labels.charAt(u) - 'a'] = 1;
        for (int v : tree[u]) {
            if (v == p) continue;
            int[] child = dfs(v, u, tree, labels, ans);
            for (int i = 0; i < 26; i++) freq[i] += child[i];
        }
        ans[u] = freq[labels.charAt(u) - 'a'];
        return freq;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

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

func countSubTrees(n int, edges [][]int, labels string) []int {
    tree := make([][]int, n)
    for _, e := range edges {
        tree[e[0]] = append(tree[e[0]], e[1])
        tree[e[1]] = append(tree[e[1]], e[0])
    }
    ans := make([]int, n)
    var dfs func(u, p int) [26]int
    dfs = func(u, p int) [26]int {
        freq := [26]int{}
        freq[labels[u]-'a'] = 1
        for _, v := range tree[u] {
            if v == p {
                continue
            }
            child := dfs(v, u)
            for i := 0; i < 26; i++ {
                freq[i] += child[i]
            }
        }
        ans[u] = freq[labels[u]-'a']
        return freq
    }
    dfs(0, -1)
    return ans
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class Solution:
    def countSubTrees(self, n: int, edges: list[list[int]], labels: str) -> list[int]:
        tree = [[] for _ in range(n)]
        for u, v in edges:
            tree[u].append(v)
            tree[v].append(u)
        ans = [0] * n
        def dfs(u: int, p: int) -> list[int]:
            freq = [0] * 26
            freq[ord(labels[u]) - ord('a')] = 1
            for v in tree[u]:
                if v == p:
                    continue
                child = dfs(v, u)
                for i in range(26):
                    freq[i] += child[i]
            ans[u] = freq[ord(labels[u]) - ord('a')]
            return freq
        dfs(0, -1)
        return ans
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

impl Solution {
    pub fn count_sub_trees(n: i32, edges: Vec<Vec<i32>>, labels: String) -> Vec<i32> {
        let n = n as usize;
        let mut tree = vec![vec![]; n];
        for e in edges.iter() {
            tree[e[0] as usize].push(e[1] as usize);
            tree[e[1] as usize].push(e[0] as usize);
        }
        let mut ans = vec![0; n];
        fn dfs(u: usize, p: usize, tree: &Vec<Vec<usize>>, labels: &[u8], ans: &mut Vec<i32>) -> [i32; 26] {
            let mut freq = [0; 26];
            freq[(labels[u] - b'a') as usize] = 1;
            for &v in &tree[u] {
                if v == p { continue; }
                let child = dfs(v, u, tree, labels, ans);
                for i in 0..26 {
                    freq[i] += child[i];
                }
            }
            ans[u] = freq[(labels[u] - b'a') as usize];
            freq
        }
        dfs(0, n, &tree, labels.as_bytes(), &mut ans);
        ans
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

class Solution {
    countSubTrees(n: number, edges: number[][], labels: string): number[] {
        const tree: number[][] = Array.from({length: n}, () => []);
        for (const [u, v] of edges) {
            tree[u].push(v);
            tree[v].push(u);
        }
        const ans = Array(n).fill(0);
        function dfs(u: number, p: number): number[] {
            const freq = Array(26).fill(0);
            freq[labels.charCodeAt(u) - 97] = 1;
            for (const v of tree[u]) {
                if (v === p) continue;
                const child = dfs(v, u);
                for (let i = 0; i < 26; i++) freq[i] += child[i];
            }
            ans[u] = freq[labels.charCodeAt(u) - 97];
            return freq;
        }
        dfs(0, -1);
        return ans;
    }
}

Complexity

  • ⏰ Time complexity: O(n), where n is the number of nodes, since each node is visited once and label merging is constant time.
  • 🧺 Space complexity: O(n), for the tree and recursion stack.