【Leetcode】2416. Sum of Prefix Scores of Strings

题目地址：

https://leetcode.com/problems/sum-of-prefix-scores-of-strings/

给定若干非空字符串组成的长$n$数组$A$，字符串只含小写英文字母，要求返回一个数组$C$，使得$C[i]={\sum }_{k}f(A[i][0:k])$，$f(s)$指的是$A$中以$s$为前缀的字符串个数。

直接用Trie，在插入的时候将经过的节点的计数都加$1$，这样$C[i]$即为$A[i]$途径节点的计数之总和。代码如下：

class Solution {
 public:
  struct Node {
    Node* ne[26];
    int cnt;
    Node() : cnt(0) { fill(ne, ne + 26, nullptr); }
  };

  Node* root;

  void insert(string& s) {
    auto p = root;
    for (int i = 0; i < s.size(); i++) {
      int id = s[i] - 'a';
      if (!p->ne[id]) p->ne[id] = new Node();
      p = p->ne[id];
      p->cnt++;
    }
  }

  vector<int> sumPrefixScores(vector<string>& words) {
    root = new Node();
    for (auto& s : words) insert(s);
    vector<int> res;
    for (auto& s : words) {
      auto p = root;
      int c = 0;
      for (int i = 0; i < s.size(); i++) {
        p = p->ne[s[i] - 'a'];
        c += p->cnt;
      }

      res.push_back(c);
    }

    return res;
  }
};

时空复杂度$O({\sum }_i{l}_{A[i]})$。