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

题目地址：

https://leetcode.com/problems/number-of-nodes-in-the-sub-tree-with-the-same-label/description/

给定一棵$n$阶无向树，顶点编号$0$开始。每个顶点有个字符标记，以字符串$s$给出，$s[u]$是点$u$的标记。要求返回一个长$n$数组$a$使得$a[u]$是$u$子树里$s[u]$出现的次数。

思路是DFS。代码如下：

class Solution {
public:
  vector<int> countSubTrees(int n, vector<vector<int>> &es, string &ls) {
    vector<int> h(n, -1), e(n - 1 << 1), ne(n - 1 << 1);
    int idx = 0;
    auto add = [&](int a, int b) { e[idx] = b, ne[idx] = h[a], h[a] = idx++; };
    for (auto &e : es) {
      int a = e[0], b = e[1];
      add(a, b), add(b, a);
    }

    vector<int> res(n);
    auto dfs = [&](this auto &&dfs, int u, int from) -> vector<int> {
      vector<int> cnt(26);
      for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v != from) {
          auto c = dfs(v, u);
          for (int i = 0; i < 26; i++)
            cnt[i] += c[i];
        }
      }

      res[u] = ++cnt[ls[u] - 'a'];
      return cnt;
    };
    dfs(0, -1);
    return res;
  }
};

时空复杂度$O(n)$。