HDU 2222 Keywords Search

HDU 2222 Keywords Search

#include  < iostream >
#include  < queue >

using   namespace  std;

struct  Trie {
    int    cnt;
    Trie*  next[26], *prefix;
} Table[ 500000 ];

int    index =   0 ;
Trie *  root;
char   str[ 1000001 ];

void  init() {
    index= 0;
    root= &Table[index];
    memset( Table[0].next, 0, sizeof(Table[0].next) );
    Table[0].cnt= 0; Table[0].prefix= NULL;
}

void  insert(  char *  s )
{
    Trie* r= root;
    while( *s ){
        int t= *s- 'a';
        if( !r->next[t] ){
            ++index;
            memset( Table[index].next, 0, sizeof(Table[index].next) );
            Table[index].cnt= 0; Table[index].prefix= NULL;
            r->next[t]= &Table[index];
        }
        r= r->next[t]; s++;
    }    
    r->cnt++;
}

void  get_prefix()
{
    Trie* r= root;
    queue<Trie*> q;
    q.push( root ); root->prefix= NULL;
    
    while( !q.empty() ){
        Trie* father= q.front(); q.pop();
        for( int i= 0; i< 26; ++i )
        if( father->next[i] ){
            Trie* tmp= father->prefix;
            while( tmp && !tmp->next[i] ) tmp= tmp->prefix;
            if( !tmp ) father->next[i]->prefix= root;
            else       father->next[i]->prefix= tmp->next[i];
            
            q.push( father->next[i] );
        }
    }
}

int  ac_run() {
    Trie* p= root;
    int ans= 0;
    char* s= str;
    while( *s ){
        int t= *s- 'a';
        while( !p->next[t] && p!= root ) p= p->prefix;
        p= p->next[t];
        if( !p ) p= root;
        Trie* tp= p;
        while( tp!= root && tp->cnt!= -1 ){
            ans+= tp->cnt; 
            tp->cnt= -1; 
            tp= tp->prefix;}
        s++;
    }
    return ans;
}

int  main()
{
    int test;
    scanf("%d",&test );
    while( test-- )
    {
        int n; char s[55];
        scanf("%d",&n ); init(); getchar();
        for( int i= 0; i< n; ++i ){
            gets(s); insert(s); }
        gets(str); get_prefix();
        printf("%d\n", ac_run() );    
    }
    
    return 0;
}