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

• Time: $O(n)$
• Space: $O(n)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 class Solution { public: vector countSubTrees(int n, vector>& edges, string labels) { vector ans(n); vector> graph(n); for (const vector& e : edges) { const int u = e[0]; const int v = e[1]; graph[u].push_back(v); graph[v].push_back(u); } dfs(graph, 0, -1, labels, ans); return ans; } private: vector dfs(const vector>& graph, int u, int parent, const string& labels, vector& ans) { vector count(26); // count of letters down from this u for (const int v : graph[u]) { if (v == parent) continue; vector childCount = dfs(graph, v, u, labels, ans); for (int i = 0; i < 26; ++i) count[i] += childCount[i]; } ans[u] = ++count[labels[u] - 'a']; // the u itself return count; } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 class Solution { public int[] countSubTrees(int n, int[][] edges, String labels) { int[] ans = new int[n]; List[] graph = new List[n]; for (int i = 0; i < n; ++i) graph[i] = new ArrayList<>(); for (int[] e : edges) { final int u = e[0]; final int v = e[1]; graph[u].add(v); graph[v].add(u); } dfs(graph, 0, -1, labels, ans); return ans; } private int[] dfs(List[] graph, int u, int parent, final String labels, int[] ans) { int[] count = new int[26]; // count of letters down from this u for (final int v : graph[u]) { if (v == parent) continue; int[] childCount = dfs(graph, v, u, labels, ans); for (int i = 0; i < 26; ++i) count[i] += childCount[i]; } ans[u] = ++count[labels.charAt(u) - 'a']; // the u itself return count; } }