hdu 1251 统计难题（求前缀出现了多少次）

题目链接：http://acm.hdu.edu.cn/showproblem.php?pid=1251

#include <iostream>
#include <cstring>
#include <cstdio>
#define mem(a,b) memset(a,b,sizeof(a));
using namespace std;
typedef long long ll;
const int maxn = 500005;
const ll INF = 0x3f3f3f3f;
int tot,n,trie[maxn][26],sum[maxn];
bool vis[maxn],flag;
void Insert(char *s,int rt)
{
    int len = strlen(s);
    for(int i = 0; i < len; i++)
    {
        int k = s[i] -'a';
        if(trie[rt][k] == 0) trie[rt][k] = ++tot;
        sum[trie[rt][k]]++;
        rt = trie[rt][k];
    }
}
int Find(char *s,int rt)
{
    int len = strlen(s);
    for(int i = 0; i < len; i++)
    {
        int k = s[i] -'a';
        rt = trie[rt][k];
        if(rt == 0) return 0;
    }

    return sum[rt];
}
int main()
{
    tot = 0;
    int rt = 0;
    flag = false;
    char s[105];
    while(1)
    {
        gets(s);
        if(s[0] == NULL)
        {
            break;
        }
            Insert(s,rt);
    }
    while(cin >> s) {
            cout << Find(s,rt) << endl;
    }
    return 0;
}