[hihoCoder] Trie树

This is a application of the Trie data structure, with minor extension. The critical part in this problem is to count all the words that have a particualr prefix and the problem has given nice hints to make this extension.

Two functions require to be implemented: add a word to the Trie and search the Trie for all words with a particular prefix.

The code is as follows.

If you are not familiar with Trie, you may refer to this solution first to get the basic idea of it.

 1 #include <iostream>
 2 
 3 using namespace std;
 4 
 5 class TrieNode {
 6 public:
 7     int count;
 8     TrieNode* children[26];
 9     TrieNode() {
10         count = 0;
11         for (int i = 0; i < 26; i++)
12             children[i] = NULL;
13     }
14 };
15 
16 class Dictionary {
17 public:
18     Dictionary() {
19         root = new TrieNode();
20     }
21     
22     void insert(char* word) {
23         TrieNode* run = root;
24         for (int i = 0; word[i]; i++) {
25             if (!(run -> children[word[i] - 'a']))
26                 run -> children[word[i] - 'a'] = new TrieNode();
27             run = run -> children[word[i] - 'a'];
28             run -> count++;
29         }
30     }
31     
32     int search(char* prefix) {
33         TrieNode* run = root;
34         for (int i = 0; prefix[i]; i++) {
35             if (run)
36                 run = run -> children[prefix[i] - 'a'];
37             else break;
38         }
39         if (!run) return false;
40         return run -> count;
41     }
42     
43 private:
44     TrieNode* root;
45 };
46 
47 int main(void) {
48     int dictSize;
49     while (scanf("%d", &dictSize) != EOF) {
50         Dictionary dictionary;
51         char word[20];
52         for (int i = 0; i < dictSize; i++) {
53             scanf("%s", word);
54             dictionary.insert(word);
55         }
56         int querySize;
57         scanf("%d", &querySize);
58         char prefix[20];
59         for (int i = 0; i < querySize; i++) {
60             scanf("%s", prefix);
61             printf("%d\n", dictionary.search(prefix));
62         }
63     }
64     return 0;
65 }