第十二章Trie树（字典树）解决HDU1251

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

#include <iostream>
using namespace std;

struct TrieNode
{
    TrieNode *children[26];
    bool flag;
    int cnt;//前缀数目
};

void InitTrieNode(TrieNode *tn)
{
    if(!tn)
    {
        exit(1);
    }
    for(int i=0;i<26;i++)
    {
        (tn->children)[i]=NULL;
    }
    tn->flag=false;
    tn->cnt=0;
}

//把长度为n的字符串a插入Trie树
void TrieTreeInsert(TrieNode **root,char *a,int n)
{
    if(!(*root))
    {
        *root=new TrieNode;
        InitTrieNode(*root);
    }
    TrieNode *tn=*root;
    
    for(int i=0;i<n;i++)
    {
        TrieNode *child=(tn->children)[a[i]-'a'];
        if(!child)
        {
            child=new TrieNode;
            InitTrieNode(child);
            (tn->children)[a[i]-'a']=child;
        }
        if(child->flag==false)
        {
            child->flag=true;
            child->cnt++;
        }
        else
        {
            child->cnt++;
        }
        tn=child;
    }
}

//统计以字符串b为前缀的单词数量
int Count(TrieNode *root,char *b,int n)
{
    int cnt=0;
    if(root)
    {
        TrieNode *tn=root;
        for(int i=0;i<n;i++)
        {
            TrieNode *child=(tn->children)[b[i]-'a'];
            if(child)
            {
                tn=child;
            }
            else
            {
                tn=NULL;
                break;
            }
        }
        if(tn!=root && tn!=NULL)
        {
            cnt=tn->cnt;
        }
    }
    return cnt;
}

int main()
{
    char a[10];
    TrieNode *root=NULL;
    while(gets(a))
    {
        int len=strlen(a);
        if(len==0)
        {
            break;
        }
        TrieTreeInsert(&root,a,len);
    }
   
    char b[10];
    while(gets(b))
    {
        int len=strlen(b);
        cout<<Count(root,b,len)<<endl;
    }
    return 0;
}