Competitive-Programming/Data_Structures/Trie_Strings.cpp at main · MohammadMazrouei/Competitive-Programming · 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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#include <bits/stdc++.h>
using namespace std;

// Trie on Strings, O(n)
struct Trie {
    int sz;
    vector<int> cnt_prefix, cnt_word;
    vector<array<int, 26>> nxt;

    Trie(int n) : sz(1) {
        cnt_prefix.assign(n, 0);
        cnt_word.assign(n, 0);
        nxt.assign(n, array<int, 26>());
    }

    void insert(const string &s) {
        int v = 0;
        for (auto ch : s) {
            int ind = ch - 'a';
            if (!nxt[v][ind]) {
                nxt[v][ind] = sz++;
            }
            v = nxt[v][ind];
            cnt_prefix[v]++;
        }
        cnt_word[v]++;
    }

    void erase(const string &s) {
        int v = 0;
        for (auto ch : s) {
            int ind = ch - 'a';
            v = nxt[v][ind];
            cnt_prefix[v]--;
            assert(cnt_prefix[v] >= 0);
        }
        cnt_word[v]--;
        assert(cnt_word[v] >= 0);
    }

    int count_prefix(const string &s) {
        int v = 0;
        for (auto ch : s) {
            int ind = ch - 'a';
            if (!nxt[v][ind]) {
                return 0;
            }
            v = nxt[v][ind];
        }
        return cnt_prefix[v];
    }

    int count_word(const string &s) {
        int v = 0;
        for (auto ch : s) {
            int ind = ch - 'a';
            if (!nxt[v][ind]) {
                return 0;
            }
            v = nxt[v][ind];
        }
        return cnt_word[v];
    }
};

void solve() {
    int q;
    cin >> q;

    const int N = 1e6;
    Trie trie(N);
    while (q--) {
        int op;
        string s;
        cin >> op >> s;
        if (op == 1) {
            trie.insert(s);
        }
        else if (op == 2) {
            trie.erase(s);
        }
        else if (op == 3) {
            cout << trie.count_prefix(s) << '\n';
        }
        else if (op == 4) {
            cout << trie.count_word(s) << '\n';
        }
    }
}

int32_t main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int t = 1;
    //cin >> t;

    while (t--) {
        solve();
    }

    return 0;
}