hdu 1251 统计难题 Trie树

#include <iostream>
#include <cctype>
#include <cstring>
#include <string>
#include <cmath>
#include <algorithm>
#include <vector>
#include <map>
#include <set>
#include <cstdio>
#include <cstdlib>
#define inf 0x3f3f3f3f
using namespace std;

struct trie{
	int count;
	trie* node[26];
	int exit;
};

trie*  creat_trie(){
	trie* temp = (trie*) malloc (sizeof(trie));
	temp->count = 0;
	temp->exit = 0;
	memset(temp->node, 0, sizeof(temp->node));
	return temp;
}

void insert(trie* root ,string str){
	int n = str.size();
	trie* nod = root;
	for(int i = 0; i < n; i++){
		int id = str[i] - 'a';
		if(nod->node[id] == NULL){
			nod->node[id] = creat_trie();
		}
		nod = nod->node[id];
		nod->count ++;
	}
}
		
int query(trie* root, string str){
	int n = str.size();
	trie* nod = root;
	for(int i = 0; i < n; i++){
		nod = nod->node[str[i] - 'a'];
		if(nod == NULL)	return 0;
	}
	return nod->count;
}
		
	
int main(){
	//freopen("ztest.txt","r",stdin);
	trie* root = creat_trie();
	string str;
	int flag = 1;
	while(getline(cin, str)){
		if(!flag){
			printf("%d\n",query(root, str));
		}
		else{
			if(str.size()){
				insert(root, str);
			}
			else{
				flag = 0;
			}
		}
	}
	return 0;
}