HDU 2222 Keywords Search ( AC 自动机 )

#include <iostream>
#include <cstdio>
#include <cstring>
using namespace std;
#define CHILD_SIZE 26

struct TrieNode {
    TrieNode* fail_pointer;
    TrieNode* childs[CHILD_SIZE];
    int occurrences;
    TrieNode() {
        fail_pointer = NULL;
        occurrences = 0;
        for ( int i = 0; i < CHILD_SIZE; ++i ) {
            childs[i] = NULL;
        }
    }
};

TrieNode* root;
int head_index = 0;
int tail_index = 0;
TrieNode* Q[500010];
char text[1000010];
char keyword[51];

void insert_string( char* str ) {
    TrieNode* temp = root;
    const int str_len = strlen( str );
    for ( int i = 0; i < str_len; ++i ) {
        const int index = str[i] - 'a';
        if ( !temp->childs[index] )
            temp->childs[index] = new TrieNode();
        temp = temp->childs[index];
    }
    temp->occurrences++;
}

void build_ac_automata() {
    Q[tail_index++] = root;
    while ( tail_index != head_index ) {
        TrieNode* parent = Q[head_index++];
        for ( int i = 0; i < CHILD_SIZE; ++i ) {
            if ( !parent->childs[i] ) continue;
            if ( parent == root ) {
                parent->childs[i]->fail_pointer = root;
            } else {
                TrieNode* moving = parent->fail_pointer;
                while ( moving ) {
                    if ( moving->childs[i] ) {
                        parent->childs[i]->fail_pointer = moving->childs[i];
                        break;
                    }
                    moving = moving->fail_pointer;
                }
                if ( !moving ) {
                    parent->childs[i]->fail_pointer = root;
                }
            }
            Q[tail_index++] = parent->childs[i];
        }
    }
}

int query_string() {

    int res = 0;
    TrieNode* node = root;
    const int str_len = strlen( text );

    for ( int i = 0; i < str_len; ++i ) {
        const int index = text[i] - 'a';
        while ( node->childs[index] == NULL && node != root ) {
            node = node->fail_pointer;
        }
        node = node->childs[index];
        if ( !node ) {
            node = root;
        }
        TrieNode* temp = node;
        while ( temp != root && temp->occurrences != -1 ) {
            res += temp->occurrences;
            temp->occurrences = -1;
            temp = temp->fail_pointer;
        }
    }

    return res;

}

int main() {

    int test_num;
    int string_num;
    scanf( "%d", &test_num );

    while ( test_num-- ) {

        head_index = 0;
        tail_index = 0;
        root = new TrieNode();

        scanf( "%d", &string_num );
        getchar();

        for ( int i = 0; i < string_num; ++i ) {
            gets( keyword );
            insert_string( keyword );
        }

        build_ac_automata();
        scanf( "%s", text );
        int res = query_string();

        printf( "%d\n", res );

    }

    return 0;

}