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

medium tree dfs bfs counting

Problem

Given a rooted tree of n nodes labeled with lowercase letters, return an array where ans[i] is the number of nodes in the subtree rooted at i that share node i's label.

Inputn = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], labels = "abaedcd"

Output[2,1,1,1,1,1,1]

Node 0 has label 'a' and its subtree contains 2 'a's (nodes 0 and 2).

edges (JSON)

labels

def count_subtrees(n, edges, labels):
    g = [[] for _ in range(n)]
    for u, v in edges:
        g[u].append(v); g[v].append(u)
    ans = [0] * n
    def dfs(u, parent):
        cnt = [0] * 26
        cnt[ord(labels[u]) - 97] = 1
        for v in g[u]:
            if v == parent: continue
            sub = dfs(v, u)
            for i in range(26):
                cnt[i] += sub[i]
        ans[u] = cnt[ord(labels[u]) - 97]
        return cnt
    dfs(0, -1)
    return ans

function countSubTrees(n, edges, labels) {
  const g = Array.from({ length: n }, () => []);
  for (const [u, v] of edges) { g[u].push(v); g[v].push(u); }
  const ans = new Array(n).fill(0);
  function dfs(u, parent) {
    const cnt = new Array(26).fill(0);
    cnt[labels.charCodeAt(u) - 97] = 1;
    for (const v of g[u]) {
      if (v === parent) continue;
      const sub = dfs(v, u);
      for (let i = 0; i < 26; i++) cnt[i] += sub[i];
    }
    ans[u] = cnt[labels.charCodeAt(u) - 97];
    return cnt;
  }
  dfs(0, -1);
  return ans;
}

class Solution {
    int[] ans;
    List<List<Integer>> g;
    String labels;
    public int[] countSubTrees(int n, int[][] edges, String labels) {
        this.labels = labels;
        ans = new int[n];
        g = new ArrayList<>();
        for (int i = 0; i < n; i++) g.add(new ArrayList<>());
        for (int[] e : edges) { g.get(e[0]).add(e[1]); g.get(e[1]).add(e[0]); }
        dfs(0, -1);
        return ans;
    }
    int[] dfs(int u, int parent) {
        int[] cnt = new int[26];
        cnt[labels.charAt(u) - 'a'] = 1;
        for (int v : g.get(u)) {
            if (v == parent) continue;
            int[] sub = dfs(v, u);
            for (int i = 0; i < 26; i++) cnt[i] += sub[i];
        }
        ans[u] = cnt[labels.charAt(u) - 'a'];
        return cnt;
    }
}

vector<int> ans;
vector<vector<int>> g;
string labels;
vector<int> dfs(int u, int parent) {
    vector<int> cnt(26, 0);
    cnt[labels[u] - 'a'] = 1;
    for (int v : g[u]) {
        if (v == parent) continue;
        auto sub = dfs(v, u);
        for (int i = 0; i < 26; i++) cnt[i] += sub[i];
    }
    ans[u] = cnt[labels[u] - 'a'];
    return cnt;
}
vector<int> countSubTrees(int n, vector<vector<int>>& edges, string lbl) {
    labels = lbl; ans.assign(n, 0); g.assign(n, {});
    for (auto& e : edges) { g[e[0]].push_back(e[1]); g[e[1]].push_back(e[0]); }
    dfs(0, -1);
    return ans;
}

Explanation

For each node we want to know how many nodes in its subtree share its letter. The neat idea is to have every subtree report back a tally of all 26 letters it contains, and let parents add up their children's tallies.

We first build an adjacency list from the edges, then DFS from node 0 (passing a parent so we never go back up). Each dfs(u) starts a count array cnt of size 26, marks its own label, and then folds in every child's count array.

Once a node has merged all of its children's tallies, the answer for that node is simply cnt[label of u] — how many times its own letter appears in the whole subtree, including itself. The node then returns cnt upward so its parent can keep accumulating.

Example: with labels "abaedcd", node 0 has label 'a'. Its subtree contains nodes 0 and 2, both 'a', so ans[0] = 2. Every other node is the only one of its letter in its subtree, giving [2,1,1,1,1,1,1].

Because each node merges a fixed-size 26-slot array once, the work per node is constant, making the whole sweep efficient.

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