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

Med

#1400Time: 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).1 company

Patterns

Counting

Algorithms

Depth-First Search Breadth-First Search

Data structures

Hash Table Tree

Companies

Samsung

Prompt

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] = [a_i, b_i], which means there is an edge between nodes a_i and b_i in the tree.

Return an array of size n where ans[i] is the number of nodes in the subtree of the i^th 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 <= 10⁵
edges.length == n - 1
edges[i].length == 2
0 <= a_i, b_i < n
a_i != b_i
labels.length == n
labels is 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 edges array.
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 ans of size n.
Iterate through each node i from 0 to n-1.
For each i, perform a traversal (like BFS or DFS) starting from i and 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 ans array.

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.

1import java.util.*;2 3class Solution {4    public int[] countSubTrees(int n, int[][] edges, String labels) {5        // Build adjacency list for the directed tree6        List<List<Integer>> adj = new ArrayList<>(n);7        for (int i = 0; i < n; i++) {8            adj.add(new ArrayList<>());9        }10        // We need to know parent-child relationships.11        // Let's build an undirected graph first, then do a BFS to establish parent-child.12        List<List<Integer>> undirectedAdj = new ArrayList<>(n);13        for (int i = 0; i < n; i++) {14            undirectedAdj.add(new ArrayList<>());15        }16        for (int[] edge : edges) {17            undirectedAdj.get(edge[0]).add(edge[1]);18            undirectedAdj.get(edge[1]).add(edge[0]);19        }20 21        Queue<Integer> q = new LinkedList<>();22        q.offer(0);23        boolean[] visited = new boolean[n];24        visited[0] = true;25 26        while (!q.isEmpty()) {27            int u = q.poll();28            for (int v : undirectedAdj.get(u)) {29                if (!visited[v]) {30                    visited[v] = true;31                    adj.get(u).add(v); // Add directed edge from parent u to child v32                    q.offer(v);33                }34            }35        }36 37        int[] ans = new int[n];38        for (int i = 0; i < n; i++) {39            char targetLabel = labels.charAt(i);40            int count = 0;41            // Traverse subtree of i42            Queue<Integer> subtreeQueue = new LinkedList<>();43            subtreeQueue.offer(i);44            while (!subtreeQueue.isEmpty()) {45                int curr = subtreeQueue.poll();46                if (labels.charAt(curr) == targetLabel) {47                    count++;48                }49                for (int child : adj.get(curr)) {50                    subtreeQueue.offer(child);51                }52            }53            ans[i] = count;54        }55        return ans;56    }57}

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).

This efficient approach avoids re-computation by using a single Depth-First Search (DFS) traversal from the root of the tree. It works in a post-order fashion: for any node, it first recursively processes all its children and then computes its own result. Each recursive call returns the frequency of all labels ('a' through 'z') in the subtree it just processed. The parent node then aggregates these frequencies from all its children to determine the frequencies for its own subtree.

Algorithm

Build an undirected adjacency list representation of the tree from the edges.
Initialize an answer array ans of size n.
Define a recursive DFS function that takes the current node, its parent, the adjacency list, labels, and the ans array as arguments. This function will return an array of size 26 with label counts for the current node's subtree.
Start the traversal from the root: dfs(0, -1, ...).
In the DFS function for a node:
- Create a frequency array counts of size 26 for the node's subtree. Initialize the count for the node's own label to 1.
- For each neighbor of the node (if it's not the parent):
  - Recursively call DFS on the neighbor.
  - Add the returned frequency array from the child's subtree to the current node's counts array.
- After visiting all children, the counts array holds the total frequencies for the node's subtree.
- Set ans[node] to counts[labels.charAt(node) - 'a'].
- Return the counts array.
After the initial DFS call completes, return the ans array.

Walkthrough

First, we build an adjacency list for the undirected graph from the edges array. We don't need to build a directed tree explicitly; we can handle parent-child relationships during the DFS traversal by passing the parent node as an argument. We define a recursive DFS function, say dfs(node, parent, adj, labels, ans). This function will traverse the subtree rooted at node and return an array of size 26, representing the counts of each character label in that subtree. The ans array, which will store the final result, is passed by reference to the DFS function. Inside the dfs(node, parent, ...) function:

Create a new integer array nodeCounts of size 26 to store label frequencies for the subtree rooted at node.
Increment the count for the label of the current node: nodeCounts[labels.charAt(node) - 'a'] = 1.
Iterate through all neighbors of node in the adjacency list.
If a neighbor is the parent of the current node, skip it to avoid going up the tree.
For each child neighbor, make a recursive call: childCounts = dfs(neighbor, node, ...).
The returned childCounts array contains label frequencies for the child's subtree. Merge these counts into nodeCounts by adding them element-wise: for (int i = 0; i < 26; i++) { nodeCounts[i] += childCounts[i]; }.
After iterating through all children, nodeCounts now contains the complete label frequencies for the subtree rooted at node.
The answer for the current node is the count of its own label in its subtree, which is nodeCounts[labels.charAt(node) - 'a']. Store this in ans[node].
Finally, return the nodeCounts array to the caller (the parent of node). The process starts by calling dfs(0, -1, adj, labels, ans). The -1 indicates that the root (node 0) has no parent.

1import java.util.*;2 3class Solution {4    public int[] countSubTrees(int n, int[][] edges, String labels) {5        List<List<Integer>> adj = new ArrayList<>(n);6        for (int i = 0; i < n; i++) {7            adj.add(new ArrayList<>());8        }9        for (int[] edge : edges) {10            adj.get(edge[0]).add(edge[1]);11            adj.get(edge[1]).add(edge[0]);12        }13 14        int[] ans = new int[n];15        dfs(0, -1, adj, labels, ans);16        return ans;17    }18 19    private int[] dfs(int node, int parent, List<List<Integer>> adj, String labels, int[] ans) {20        // Store counts of all characters in the subtree of the current node.21        int[] nodeCounts = new int[26];22        char label = labels.charAt(node);23        nodeCounts[label - 'a'] = 1;24 25        for (int neighbor : adj.get(node)) {26            if (neighbor == parent) {27                continue;28            }29            // Get counts from the child's subtree.30            int[] childCounts = dfs(neighbor, node, adj, labels, ans);31            // Add child's counts to the current node's counts.32            for (int i = 0; i < 26; i++) {33                nodeCounts[i] += childCounts[i];34            }35        }36 37        // The answer for the current node is the count of its own label.38        ans[node] = nodeCounts[label - 'a'];39        return nodeCounts;40    }41}

Complexity

Time

O(N). Each node is visited once. At each node, we perform a constant number of operations plus merging count arrays from its children. The merging operation takes O(26) time. Since each child's count array is merged exactly once (into its parent's), the total time for all merge operations across the entire tree is O(N * 26). Thus, the total time complexity is O(N).

Space

O(N). O(N) for the adjacency list. The recursion stack can go up to O(H) deep, where H is the height of the tree. In the worst case (a skewed tree), H=N, so the stack depth is O(N). At each level of recursion, we create a new count array of size 26. So, the space used by the recursion stack and the count arrays is O(H * 26), which is O(N) in the worst case. The `ans` array also takes O(N) space. Total space is O(N).

Trade-offs

Pros

Very efficient, solving the problem in a single pass over the tree.
Optimal time complexity for this problem.

Cons

The logic is slightly more complex than the brute-force approach, involving recursion and passing data up the call stack.
The space complexity for the recursion stack can be O(N) in the worst case for a skewed tree.

Solutions

Solution

1class Solution {2private3  List<Integer>[] g;4private5  String labels;6private7  int[] ans;8private9  int[] cnt;10public11  int[] countSubTrees(int n, int[][] edges, String labels) {12    g = new List[n];13    Arrays.setAll(g, k->new ArrayList<>());14    for (int[] e : edges) {15      int a = e[0], b = e[1];16      g[a].add(b);17      g[b].add(a);18    }19    this.labels = labels;20    ans = new int[n];21    cnt = new int[26];22    dfs(0, -1);23    return ans;24  }25private26  void dfs(int i, int fa) {27    int k = labels.charAt(i) - 'a';28    ans[i] -= cnt[k];29    cnt[k]++;30    for (int j : g[i]) {31      if (j != fa) {32        dfs(j, i);33      }34    }35    ans[i] += cnt[k];36  }37}

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

1class Solution:2    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]: def dfs(i, fa): ans[i] -= cnt[labels[i]] cnt[labels[i]] += 1 for j in g[i]: if j != fa: dfs(j, i) ans[i] += cnt[labels[i]] g = defaultdict(list) for a, b in edges: g[a]. append(b) g[b]. append(a) cnt = Counter() ans = [0] * n dfs(0, - 1) return ans

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.

← Previous

Previous problem

Next problem