Number of Nodes in the Sub-Tree With the Same Label
MedPrompt
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 <= 105edges.length == n - 1edges[i].length == 20 <= ai, bi < nai != bilabels.length == nlabelsis consisting of only of lowercase English letters.
Approaches
2 approaches with complexity analysis and trade-offs.
This approach iterates through every node in the tree. For each node, it performs a separate traversal (like DFS or BFS) to visit all nodes in its subtree. During this traversal, it counts the nodes that have the same label as the starting node of the subtree.
Algorithm
- Build an undirected adjacency list from the
edgesarray. - Perform a Breadth-First Search (BFS) starting from the root (node 0) to build a directed adjacency list representing parent-to-child relationships. This defines the subtrees.
- Initialize an answer array
ansof sizen. - Iterate through each node
ifrom0ton-1. - For each
i, perform a traversal (like BFS or DFS) starting fromiand visiting only its descendants. - During this traversal, count how many nodes have the same label as node
i. - Store this count in
ans[i]. - Return the
ansarray.
Walkthrough
First, we need to represent the tree in a way that's easy to traverse, like an adjacency list. Since the input edges are undirected, we build an undirected graph representation.
The problem defines a rooted tree at node 0. To correctly identify subtrees, we need parent-child relationships. We can establish these with a preliminary traversal (e.g., BFS) starting from the root (node 0). This allows us to build a directed version of the tree where edges go from parent to child.
The main part of the algorithm is a loop that iterates from i = 0 to n-1.
Inside the loop, for each node i, we start a new traversal (e.g., DFS) from i. This traversal will only move from a parent to its children, effectively exploring the subtree of i.
We maintain a counter for this traversal. Whenever we visit a node j in the subtree of i, we check if labels[j] is the same as labels[i]. If they match, we increment the counter.
After the traversal for node i is complete, the value of the counter is the answer for i, so we store it in ans[i].
This process is repeated for all n nodes.
import java.util.*; class Solution { public int[] countSubTrees(int n, int[][] edges, String labels) { // Build adjacency list for the directed tree List<List<Integer>> adj = new ArrayList<>(n); for (int i = 0; i < n; i++) { adj.add(new ArrayList<>()); } // We need to know parent-child relationships. // Let's build an undirected graph first, then do a BFS to establish parent-child. List<List<Integer>> undirectedAdj = new ArrayList<>(n); for (int i = 0; i < n; i++) { undirectedAdj.add(new ArrayList<>()); } for (int[] edge : edges) { undirectedAdj.get(edge[0]).add(edge[1]); undirectedAdj.get(edge[1]).add(edge[0]); } Queue<Integer> q = new LinkedList<>(); q.offer(0); boolean[] visited = new boolean[n]; visited[0] = true; while (!q.isEmpty()) { int u = q.poll(); for (int v : undirectedAdj.get(u)) { if (!visited[v]) { visited[v] = true; adj.get(u).add(v); // Add directed edge from parent u to child v q.offer(v); } } } int[] ans = new int[n]; for (int i = 0; i < n; i++) { char targetLabel = labels.charAt(i); int count = 0; // Traverse subtree of i Queue<Integer> subtreeQueue = new LinkedList<>(); subtreeQueue.offer(i); while (!subtreeQueue.isEmpty()) { int curr = subtreeQueue.poll(); if (labels.charAt(curr) == targetLabel) { count++; } for (int child : adj.get(curr)) { subtreeQueue.offer(child); } } ans[i] = count; } return ans; }}Complexity
Time
O(N^2). Building the directed tree takes O(N). The main loop runs N times. Inside the loop, we traverse the subtree of node `i`. The size of the subtree can be up to O(N). In a skewed tree (like a path), the sum of subtree sizes is `N + (N-1) + ... + 1 = O(N^2)`. Therefore, the total time complexity is dominated by these traversals, leading to O(N^2).
Space
O(N). We use O(N) space for the adjacency lists. The queue used for BFS in both the tree-building step and the per-node traversal can hold up to O(N) nodes in the worst case (e.g., a star graph).
Trade-offs
Pros
Simple to understand and implement.
It correctly solves the problem by directly simulating the definition of a subtree count.
Cons
Highly inefficient due to redundant computations. The traversal for a parent node re-traverses the subtrees of all its children.
Will result in a "Time Limit Exceeded" (TLE) error for larger inputs as specified in the constraints (n up to 10^5).
Solutions
Solution
class Solution {private List<Integer>[] g;private String labels;private int[] ans;private int[] cnt;public int[] countSubTrees(int n, int[][] edges, String labels) { g = new List[n]; Arrays.setAll(g, k->new ArrayList<>()); for (int[] e : edges) { int a = e[0], b = e[1]; g[a].add(b); g[b].add(a); } this.labels = labels; ans = new int[n]; cnt = new int[26]; dfs(0, -1); return ans; }private void dfs(int i, int fa) { int k = labels.charAt(i) - 'a'; ans[i] -= cnt[k]; cnt[k]++; for (int j : g[i]) { if (j != fa) { dfs(j, i); } } ans[i] += cnt[k]; }}Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.