|
2 | 2 | #define ll long long |
3 | 3 | using namespace std; |
4 | 4 |
|
5 | | -struct trie_node { |
6 | | - ll word_count; |
7 | | - ll prefix_count; |
8 | | - trie_node* edges[26]; |
9 | | - |
10 | | - trie_node() { |
11 | | - this->word_count = 0; |
12 | | - this->prefix_count = 0; |
13 | | - for(int i = 0; i < 26; i++) { |
14 | | - this->edges[i] = NULL; |
15 | | - } |
16 | | - } |
| 5 | +const int MAXN = 260005; |
| 6 | + |
| 7 | +struct arr { |
| 8 | + int a[26]; |
| 9 | + arr() {} |
| 10 | + void clear() { memset(a,-1,sizeof(a)); } |
| 11 | + int& operator[](int i) { return a[i]; } |
17 | 12 | }; |
18 | 13 |
|
19 | | -void add_word(trie_node* &root, string &str, int index) { |
20 | | - if (index == str.length()) { |
21 | | - root->word_count += 1; |
22 | | - return; |
23 | | - } |
24 | | - root->prefix_count += 1; |
25 | | - int i = int(str[index] - 'a'); |
26 | | - if (root->edges[i] == NULL) { |
27 | | - root->edges[i] = new trie_node(); |
28 | | - } |
29 | | - add_word(root->edges[i], str, index+1); |
30 | | -} |
| 14 | +struct trie { |
| 15 | + int cnt, prefix_cnt[MAXN], word_cnt[MAXN]; |
| 16 | + arr to[MAXN]; |
31 | 17 |
|
32 | | -ll count_prefix(trie_node* root, string &prefix, int index) { |
33 | | - if (index == prefix.length()) { |
34 | | - return root->prefix_count; |
| 18 | + trie() { cnt = MAXN-1; } |
| 19 | + |
| 20 | + void clear() { for(int i = 0; i < cnt; i++) prefix_cnt[i] = word_cnt[i] = 0, to[i].clear(); cnt = 1; } |
| 21 | + |
| 22 | + void add(const string& s) { |
| 23 | + int u = 0; |
| 24 | + for(const char& c: s) { |
| 25 | + if(to[u][c-'a'] == -1) to[u][c-'a'] = cnt++; |
| 26 | + u = to[u][c-'a']; |
| 27 | + prefix_cnt[u]++; |
| 28 | + } |
| 29 | + word_cnt[u]++; |
35 | 30 | } |
36 | | - int i = int(prefix[index] - 'a'); |
37 | | - if (root->edges[i] == NULL) { |
38 | | - return 0; |
| 31 | + |
| 32 | + int prefix_count(const string& s) { |
| 33 | + int u = 0; |
| 34 | + for (const char& c: s) { |
| 35 | + if (to[u][c-'a'] == -1) return 0; |
| 36 | + u = to[u][c-'a']; |
| 37 | + } |
| 38 | + return prefix_cnt[u]; |
39 | 39 | } |
40 | | - return count_prefix(root->edges[i], prefix, index+1); |
41 | | -} |
| 40 | +} tr; |
42 | 41 |
|
43 | 42 | int main() { |
44 | | - trie_node* root = newtrie_node(); |
| 43 | + tr.clear(); |
45 | 44 | string s = "codechef"; |
46 | | - add_word(root, s, 0); |
| 45 | + tr.add(s); |
47 | 46 | s = "codeforces"; |
48 | | - add_word(root, s, 0); |
| 47 | + tr.add(s); |
49 | 48 | s = "youtube"; |
50 | | - add_word(root, s, 0); |
| 49 | + tr.add(s); |
51 | 50 | s = "google"; |
52 | | - add_word(root, s, 0); |
| 51 | + tr.add(s); |
53 | 52 | s = "code"; |
54 | | - cout<<count_prefix(root, s, 0)<<endl; |
| 53 | + cout<<tr.prefix_count(s)<<endl; |
55 | 54 | return 0; |
56 | 55 | } |
0 commit comments