Number of Nodes in the Sub-Tree With the Same Label — Medium Problem & Solution

A tree of n nodes is rooted at node 0, given by n - 1 undirected edges. labels[i] is the lowercase letter on node i.

Difficulty: Medium
Topics: Hash Table, Breadth-First Search, Counting, Depth-First Search, Graph, Trees
Asked at: Amazon, Google, Salesforce
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

A tree of n nodes is rooted at node 0, given by n - 1 undirected edges. labels[i] is the lowercase letter on node i.

Return an array where entry i is the number of nodes in the subtree of node i that carry the same label as node i — including node i itself.

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 is labelled `a`, and its subtree holds one more `a` (node 2).

Example 2

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

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^5
edges.length == n - 1
edges[i].length == 2
0 <= edges[i][0], edges[i][1] < n
labels.length == n
labels is made of lowercase English letters.

How to solve Number of Nodes in the Sub-Tree With the Same Label

Give each node a 26-entry letter tally of its subtree. Process nodes children-first, folding each child's tally into its parent; the answer for a node is its own letter's entry once the fold is complete.

Approach

Root the tree at 0 and record a traversal order with parents.
Seed every node's tally with 1 in its own letter.
Walk the order backwards: read the node's answer, then add its tally into its parent's.
Reading before merging is what keeps a node's tally to its own subtree.

Why it works

The alphabet's fixed size is what makes this linear: a tally is 26 numbers regardless of subtree size, so merging a child costs O(26) rather than O(subtree). Processing in reverse traversal order guarantees a node is read only after every descendant has been folded in, and before it is folded into its own parent.

Complexity

Time — O(26 · n)
Space — O(26 · n)

Pitfalls

Merging a node into its parent before reading its answer mixes in siblings.
The tree is given as undirected edges; the root must be fixed at node 0.
Recursion can reach depth n; an explicit stack is safer at 100 000 nodes.

Reference solution

Python

from typing import List

def countSubTrees(n: int, edges: List[List[int]], labels: str) -> List[int]:
    adj = [[] for _ in range(n)]
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)
    parent = [-1] * n
    order = []
    seen = [False] * n
    seen[0] = True
    stack = [0]
    while stack:
        u = stack.pop()
        order.append(u)
        for v in adj[u]:
            if seen[v]:
                continue
            seen[v] = True
            parent[v] = u
            stack.append(v)
    cnt = [[0] * 26 for _ in range(n)]
    for i in range(n):
        cnt[i][ord(labels[i]) - 97] = 1
    out = [0] * n
    for i in range(len(order) - 1, -1, -1):
        u = order[i]
        out[u] = cnt[u][ord(labels[u]) - 97]
        p = parent[u]
        if p >= 0:
            for c in range(26):
                cnt[p][c] += cnt[u][c]
    return out

JavaScript

var countSubTrees = function(n, edges, labels) {
    var i, c;
    var adj = [];
    for (i = 0; i < n; i++) adj.push([]);
    for (i = 0; i < edges.length; i++) {
        adj[edges[i][0]].push(edges[i][1]);
        adj[edges[i][1]].push(edges[i][0]);
    }
    var parent = [], seen = [];
    for (i = 0; i < n; i++) { parent.push(-1); seen.push(false); }
    var order = [];
    var stack = [0];
    seen[0] = true;
    while (stack.length > 0) {
        var u = stack.pop();
        order.push(u);
        for (i = 0; i < adj[u].length; i++) {
            var v = adj[u][i];
            if (seen[v]) continue;
            seen[v] = true;
            parent[v] = u;
            stack.push(v);
        }
    }
    var cnt = [];
    for (i = 0; i < n; i++) {
        var row = [];
        for (c = 0; c < 26; c++) row.push(0);
        row[labels.charCodeAt(i) - 97] = 1;
        cnt.push(row);
    }
    var out = [];
    for (i = 0; i < n; i++) out.push(0);
    for (i = order.length - 1; i >= 0; i--) {
        var w = order[i];
        out[w] = cnt[w][labels.charCodeAt(w) - 97];
        var p = parent[w];
        if (p >= 0) {
            for (c = 0; c < 26; c++) cnt[p][c] += cnt[w][c];
        }
    }
    return out;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 202 hash table problems · the whole catalogue