HDU 1251 TrieTree

http://acm.hdu.edu.cn/showproblem.php?pid=1251

 统计难题
Time Limit: 4000/2000 MS (Java/Others) Memory Limit: 131070/65535 K (Java/Others)
Total Submission(s): 8595 Accepted Submission(s): 3436


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>

using namespace std;

const int BranchNumber=26;

class TrieNode
{
    friend class TrieTree;

    public:
        TrieNode();
        //记录到此是否成为一个串
        bool isStr;    
        int count;

    private:
        TrieNode* next[BranchNumber];
};


TrieNode::TrieNode()
    :isStr(false),count(0)
{
    memset(next,NULL,sizeof(next));
}

class TrieTree
{
    public:
        TrieTree();
        void Insert(const char* word);
        int Search(const char* word);
        void DeleteTree(TrieNode* root);
        ~TrieTree();

    private:
        void CountByPrefix(TrieNode* t,int& count);
        TrieNode* root;
        TrieNode* GetNewNode();
};

TrieTree::TrieTree()
{
    root=GetNewNode();
}

TrieNode* TrieTree::GetNewNode()
{
    TrieNode* newNode;
    newNode=new TrieNode();
    if(newNode==NULL)
    {
        std::cerr<<"Memory allocation failure.\n";
        exit(1);
    }
    return newNode;
}

void TrieTree::Insert(const char* word)
{
    TrieNode* location=root;
    while(*word)
    {
        //tolower(*word);
        if(location->next[*word-'a']==NULL)
        {
            TrieNode* tempNode=GetNewNode();
            location->next[*word-'a']=tempNode;
        }
        //每插入一步，相当于有一个新串经过，指针要向下移动
        location=location->next[*word-'a'];
        location->count++;
        word++;
    }
    location->isStr=true;
}

int TrieTree::Search(const char* word)
{
    TrieNode* location=root;
    while(*word && location)
    {
        location=location->next[*word-'a'];
        word++;
    }

    if(location==NULL)
        return 0;
    else
        return location->count;
}


TrieTree::~TrieTree()
{
    DeleteTree(root);
}

void TrieTree::DeleteTree(TrieNode* root)
{
    if(root!=NULL)
    {
        for(int i=0;i<BranchNumber;i++)
            DeleteTree(root->next[i]);
        delete root;
    }
}


int main()
{
    char word[10];
    char ask[10];
    TrieTree trieTree;
    while(gets(word)) 
    {
        if(word[0] == '\0') break;
        trieTree.Insert(word);/*建树*/
    }
    while(cin>>ask)         
        cout << trieTree.Search(ask)<< endl;/*查单词*/
    
    return 0;
}

//End AC Code......
下面是自己写的插入时候再查找的代码

int TrieTree::Search(const char* word)
{
    TrieNode* location=root;
    while(*word && location)
    {
        location=location->next[*word-'a'];
        word++;
    }
    int count=0;
    CountByPrefix(location,count);
    return count;
}

void TrieTree::CountByPrefix(TrieNode* t,int& count)
{
    if(t!=NULL && t->isStr)
        count++;
    for(int i=0;i<BranchNumber;i++)
    {
        if(t->next[i]!=NULL)
            CountByPrefix(t->next[i],count);
    }
}

不过这样好像。。。就算改成非递归也超时了。。。