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

#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;
}