leetcode/Problems/2418.cpp at master · DimitrijeDobrota/leetcode · GitHub

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
36
37
38
39
40
41
42
43
44
45
46
class Solution {
    class Trie {
        struct Node {
            Node *children[26] = {nullptr};
            int count = 0;
        } node;

      public:
        void insert(const string &word) {
            Node *crnt = &node;

            for (const char c : word) {
                const auto idx = c - 'a';
                if (!crnt->children[idx]) crnt->children[idx] = new Node();
                crnt = crnt->children[idx];
                crnt->count++;
            }
        }

        int count(const string &word) const {
            const Node *crnt = &node;
            int res = 0;

            for (const char c : word) {
                const auto idx = c - 'a';
                crnt = crnt->children[idx];
                res += crnt->count;
            }

            return res;
        }
    };

  public:
    vector<int> sumPrefixScores(const vector<string> &words) const {
        vector<int> res;
        Trie trie;

        for (const auto &word : words)
            trie.insert(word);
        for (const auto &word : words)
            res.push_back(trie.count(word));

        return res;
    }
};