【字典树】统计难题(查询前缀出现次数)

统计难题
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> 
#define MAX 500010
using namespace std;

int trie[MAX][26];	//表示字典树最多有MAX个节点，每个节点的子节点有26个，用字母-'a'表示 
bool v[MAX];	//v[i]表示节点i是否是单词结束标值
int sum[MAX];	//sum[i]表示节点i被访问过的次数，即前缀出现次数 
int k = 0;	//用来赋值节点编号 

void insert(string s)
{
	int len = s.length();
	int rt = 0;
	for(int i = 0;i < len;i++)
	{
		int id = s[i]-'a';
		if(!trie[rt][id])	//第i个节点的子节点不存在，即走到尽头了 
		{
			trie[rt][id] = ++k;
		} 
		rt = trie[rt][id];	//顺着字典树往下走
		sum[rt]++;	//表示前缀出现次数 
	}
//	v[rt] = true;	表示该单词在rt节点结束 
}

bool find(string s)	//查找是否存在前缀s
{
	int len = s.length();
	int rt = 0;
	for(int i = 0;i < len;i++)
	{
		int id = s[i]-'a';
		if(!trie[rt][id])	//遍历中途没有找到s[i]字母
			return false;
		rt = trie[rt][id];	//继续遍历字典树 
	}
	return true;	//前缀遍历完
//	return v[rt];	表示存在单词s 
} 

int search(string s)	//查询前缀s出现次数 
{
	int len = s.length();
	int rt = 0;
	for(int i = 0;i < len;i++)
	{
		int id = s[i]-'a';
		if(!trie[rt][id]) return 0;	//中途断掉，返回0 
		rt = trie[rt][id];
	}
	return sum[rt];	//返回该节点被遍历的次数，即前缀出现次数 
}

int main()
{
	string s;
	while(1)
	{
		//注意：getline会读入'\n'然后吞掉. 因此直接getline，不需要getchar().如果直接按'\n'，则s[0]='\0'. 
		//如果前面有cin>> 就需要加上getchar(). 因为cin不会吞掉'\n'. 
		getline(cin,s);	
		if(s[0] == '\0') break;	
		insert(s);
	}
	while(cin>>s)
	{	
		int ans = search(s);
		cout<<ans<<endl;
	} 
	return 0;
}