zoj 3430 Detect the Virus

Problem:
Aho-Corasick Automation

Code:

#include <stdio.h>
#include <string.h>

#define KEYSIZE 256

char tmp[2060];
int virus[164];
int email[1600];

struct Node {
    Node* fail;
    Node* next[KEYSIZE];
    int count;
    bool visit;
    Node() {
        fail = NULL;
        count = 0;
        visit = false;
        memset(next, 0, sizeof(next));
    }
}*build_ac_automation_que[40000];
Node* query_temp_que[1000];

int map[128];

void base642char(char ch1, char ch2, char ch3, char ch4, int* out) {
    int i1 = map[ch1];
    int i2 = map[ch2];
    int i3 = map[ch3];
    int i4 = map[ch4];
    out[0] = (i1<<2)+(i2>>4);
    out[1] = ((i2%16)<<4)+(i3>>2);
    out[2] = ((i3%4)<<6)+i4;
}

void decode(char* in, int* out, int* size) {
    int i;
    int len = strlen(in);
    for( i=0; i<len-4; i+=4 ) {
    base642char(in[i], in[i+1], in[i+2], in[i+3], out);
    out += 3;
    }
    (*size) += (len-4)/4*3;
    int i1 = map[in[len-4]];
    int i2 = map[in[len-3]];

    out[0] = (i1<<2)+(i2>>4);
    (*size)++;

    int i3, i4;
    if( in[len-2] != '=' ) {
    i3 = map[in[len-2]];
    out[1] = ((i2%16)<<4) + (i3>>2);
    (*size)++;
    }
    if( in[len-1] != '=' ) {
    i4 = map[in[len-1]];
    out[2] = ((i3%4)<<6)+i4;
    (*size)++;
    }
}

void build_aho_corasick_automation(Node *root) {
    root->fail = NULL;
    int head, tail;
    head = tail = 0;
    build_ac_automation_que[tail++] = root;
    while( head < tail ) {
    Node *cur = build_ac_automation_que[head++];
    for(int i=0; i<KEYSIZE; i++) {
        if( cur->next[i] == NULL)
        continue;
        if( cur == root )
        cur->next[i]->fail = root;
        else {
        Node *p = cur->fail;
        while( p != NULL ) {
            if( p->next[i] != NULL ) {
            cur->next[i]->fail = p->next[i];
            break;
            }
            p = p->fail;
        }
        if( p == NULL )
            cur->next[i]->fail = root;
        }
        build_ac_automation_que[tail++] = cur->next[i];
    }
    }

}

int query(Node *root, int *str, int num ) {
    int i = 0, cnt = 0;
    int head, tail;
    head = tail = 0;
    Node *p = root;
    for( i=0; i<num; i++ ) {
    int index = str[i];
    while( p->next[index] == NULL && p!=root )
        p = p->fail;
    p = p->next[index];
    p = (p==NULL) ? root : p;
    Node *temp = p;
    while( temp != root && !temp->visit ) {
        cnt += temp->count;
        temp->visit = true;
        query_temp_que[tail++] = temp;
        temp = temp->fail;
    }
    }
    while( head < tail ) {
    Node* cur = query_temp_que[head++];
    cur->visit = false;
    }
    return cnt;
}

void insert(Node *root, int *str, int num ) {
    Node *p = root;
    int i=0;
    for( i=0; i<num; i++ ) {
    int index = str[i];
    if( p->next[index] == NULL ) {
        p->next[index] = new Node();
    }
    p = p->next[index];
    }
    p->count++;
}

void init()
{
    int i;
    for( i='a'; i<='z'; i++ ) {
    map[i] = i-'a'+26;
    }
    for( i='A'; i<='Z'; i++ ) {
    map[i] = i-'A';
    }
    for( i='0'; i<='9'; i++ ) {
    map[i] = i-'0'+52;
    }
    map['+'] = 62;
    map['/'] = 63;
}

int main()
{
    int n, m, i, j;
    Node *root;
    init();

    while( EOF != scanf("%d", &n)) {

    int decode_size;
    root = new Node();
    for( i=0; i<n; i++ ) {
        scanf("%s", tmp);
        decode_size = 0;
        decode(tmp, virus, &decode_size);
        insert(root, virus, decode_size);
    }

    build_aho_corasick_automation(root);

    scanf("%d", &m);
    for( i=0; i<m; i++ ) {
        scanf("%s", tmp);
        decode_size = 0;
        decode(tmp, email, &decode_size);

        int res = query(root, email, decode_size);
        printf("%d\n", res);
    }
    printf("\n");
    }

    return 0;
}