HDU1251 统计难题 trie树

题意：给定一些单词和一些前缀，问每个前缀出现多少次

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

第一次写trie树，数组开小交上去TLE调了好久

#include<bits/stdc++.h>

using namespace std;

const int MAXN = 2333666;
int s[MAXN], g[MAXN][26], n;
char str[20];

int find (char str[]) {
    int x = 1;
    for (int i = 0; str[i]; ++i) {
        char ch = str[i];
        if (g[x][ch - 'a'] == 0) return 0;
        x = g[x][ch - 'a'];
    }
    return s[x];    
}

void add (char str[]) {
    int x = 1;
    ++s[1];
    for (int i = 0; str[i]; ++i) {
        char ch = str[i];
        if (g[x][ch - 'a'] == 0) {
            g[x][ch - 'a'] = ++n;
        }
        x = g[x][ch - 'a'];
        ++s[x];
    }    
}

int main() {
    n = 1;
    while (gets(str) && str[0] != NULL) {
        add(str);
    }
    while (gets(str) && str[0] != NULL) {
        printf("%d\n", find(str));
    }    
    return 0;
}