HDU 1251 统计难题 -- 字典树

1251 ( 统计难题 )

#include <stdio.h>
#include <iostream>
#include <string.h>
#include <math.h>
#include <algorithm>
#include <time.h>
using namespace std;

const int MAX_26 = 26;
const int MAX_TREE = 390010;
#define FORi_26 for(int i=0;i<MAX_26;i++)
struct TrieTree{
    int ncount;
    char node; // 改名Char
    TrieTree *next[MAX_26]; // 这个指针指向一个地址
}Tt[MAX_TREE];
int tot = 1; // Tt[] array sub

bool Scan(char * str){
    while( *str = getchar() )
        if (*str == '\n'){
            *str = '\0';
            return true;
        }else if (*str == EOF)
            return false;
        else
            str++;
    return true;
}

void UpdataTree(TrieTree *node,char *str){ // 第一次进来的是头结点 
    if (*str == '\0'){
        Tt[0].ncount += 1; // 树种的单词总数
        return ;
    }
    FORi_26{ // 存在前缀往下递归
        TrieTree *temp = node->next[i];
        if ( temp != NULL && temp->node == *str){
            UpdataTree(temp,str+1);
            temp->ncount += 1;
            return ;
        }
    }
    // 不存在前缀往后添一个
    TrieTree *temp = &Tt[tot++]; // 从数组里面 申请一个结点
    FORi_26{
        if (node->next[i] == NULL)
        {    node->next[i] = temp;break;}
    }
    temp->ncount = 1;
    temp->node = *str;
    FORi_26{
        temp->next[i] = NULL;
    }
    UpdataTree(temp,str+1);
}
int SearchTree(TrieTree *node,char *str){
    int res = 0;
    int i=0;
    while(node->next[i] != NULL && *str !='\0'){
        TrieTree *temp = node->next[i];
        if(temp->node == *str){
            res = temp->ncount;
            node = temp; // 下一层
            i = 0;
            str++;
            continue;
        }
        i++;
    }
    if (*str!='\0')
        return 0;
    return res;
}

int main(){

    char str[12];
    while(Scan(str) && str[0] != '\0')
        UpdataTree( &Tt[0] , str );

    //there is a blank enter in the two parts

    while(Scan(str) && str[0] != '\0')
        printf("%d\n", SearchTree( &Tt[0] , str ));

    return 0;
}