Number of Nodes in the Sub-Tree With the Same Label
MediumUpdated: Aug 2, 2025
Practice on:
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:
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]
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
- Build the tree as an adjacency list from the edges.
- Define a DFS function that returns a frequency array for each subtree.
- For each node, recursively merge the frequency arrays of its children.
- Increment the frequency for the current node's label and set the answer for that node.
- Return the answer array after DFS traversal.
Code
Java
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;
}
}
C++
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;
}
};
Go
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
}
Python
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
Rust
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
}
}
TypeScript
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), wherenis 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.