Number of Nodes in the Sub-Tree With the Same Label LeetCode solution in C++

1
class Solution {
2
  vector<int> ans;
3
  vector<vector<int>> graph;
4
  string global;
5

6
 public:
7
  vector<int> getAns(int u, int parent) {
8
    ans[u] = 1;
9
    char rootChar = global[u];
10
    vector<int> toReturn = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
11
                            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
12

13
    for (int adj : graph[u]) {
14
      if (adj != parent) {
15
        vector<int> c = getAns(adj, u);
16
        for (int i = 0; i < 26; i++) {
17
          toReturn[i] += c[i];
18
        }
19
        ans[u] += c[rootChar - 'a'];
20
      }
21
    }
22
    toReturn[rootChar - 'a'] += 1;
23
    return toReturn;
24
  }
25

26
  vector<int> countSubTrees(int n, vector<vector<int>> &edges, string labels) {
27
    ans.resize(n);
28
    graph.resize(n);
29
    global = labels;
30

31
    for (vector<int> edge : edges) {
32
      graph[edge[0]].push_back(edge[1]);
33
      graph[edge[1]].push_back(edge[0]);
34
    }
35

36
    int root = 0;
37
    int parent = -1;
38
    getAns(root, parent);
39
    return ans;
40
  }
41
};