跳转至

HDU-1251 统计难题(单词查找树)

Problem Description

Ignatius最近遇到一个难题,老师交给他很多单词(只有小写字母组成,不会有重复的单词出现),现在老师要他统计出以某个字符串为前缀的单词数量(单词本身也是自己的前缀).

Input

输入数据的第一部分是一张单词表,每行一个单词,单词的长度不超过10,它们代表的是老师交给Ignatius统计的单词,一个空行代表单词表的结束.第二部分是一连串的提问,每行一个提问,每个提问都是一个字符串.

注意:本题只有一组测试数据,处理到文件结束.

Output

对于每个提问,给出以该字符串为前缀的单词的数量.

Sample Input

banana
band
bee
absolute
acm

ba
b
band
abc

Sample Output

2
3
1
0

#include <iostream>
#include <string>

using namespace std;

class TrieNode{
public:
    int wordNum;
    bool isWord; //到当前字符为止是否是一个完整的单词
    TrieNode *child[26];

    TrieNode() : wordNum(0), isWord(false) {
        for (auto & a : child) a = nullptr;
    }
};

class Trie {
public:
    /** Initialize your data structure here. */
    Trie() {
        root  = new TrieNode();
    }

    ~Trie(){
        clear();
    }

    /** Inserts a word into the trie. */
    void insert(string word) {
        TrieNode *p = root;
        ++(p -> wordNum);

        for (auto &a : word){
            int id = a - 'a';
            if (!p -> child[id]) p -> child[id] = new TrieNode();
            p = p -> child[id];
            ++(p -> wordNum);
        }

        p -> isWord = true;
    }

    /** Returns if the word is in the trie. */
    int search(string word) {
        TrieNode *p = root;

        for (auto a : word){
            int id = a - 'a';
            if (!p -> child[id]) return 0;
            p = p -> child[id];
        }

        return p -> wordNum;
    }

    void clear()
    {
        remove(root);

        for (int i = 0; i < 26; ++i)
            root -> child[i] = nullptr;
    }

private:
    TrieNode *root;

    void remove(TrieNode * &p)
    {
        for (int i = 0; i < 26; ++i){
            if (!p -> child[i]) continue;
            remove(p -> child[i]);
            delete p -> child[i];
            p -> child[i] = nullptr;
            if ( (--(p -> wordNum)) == 0 ) break;
        }
    }
};

int main()
{
    Trie trie;
    string inputStr;

    while (getline(cin, inputStr) && inputStr != ""){
        trie.insert(inputStr);
    }

    while (cin >> inputStr)
        cout << trie.search(inputStr) << endl;

    return 0;
}

字典树模板,坑点是编译选项时选g++Memory Limit Exceed,选C++则编译通过。