[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.

#include <iostream>

using namespace std;

class TrieNode {
public:
    int count;
    TrieNode* children[26];
    TrieNode() {
        count = 0;
        for (int i = 0; i < 26; i++)
            children[i] = NULL;
    }
};

class Dictionary {
public:
    Dictionary() {
        root = new TrieNode();
    }
    
    void insert(char* word) {
        TrieNode* run = root;
        for (int i = 0; word[i]; i++) {
            if (!(run -> children[word[i] - 'a']))
                run -> children[word[i] - 'a'] = new TrieNode();
            run = run -> children[word[i] - 'a'];
            run -> count++;
        }
    }
    
    int search(char* prefix) {
        TrieNode* run = root;
        for (int i = 0; prefix[i]; i++) {
            if (run)
                run = run -> children[prefix[i] - 'a'];
            else break;
        }
        if (!run) return false;
        return run -> count;
    }
    
private:
    TrieNode* root;
};

int main(void) {
    int dictSize;
    while (scanf("%d", &dictSize) != EOF) {
        Dictionary dictionary;
        char word[20];
        for (int i = 0; i < dictSize; i++) {
            scanf("%s", word);
            dictionary.insert(word);
        }
        int querySize;
        scanf("%d", &querySize);
        char prefix[20];
        for (int i = 0; i < querySize; i++) {
            scanf("%s", prefix);
            printf("%d\n", dictionary.search(prefix));
        }
    }
    return 0;
}