屏蔽字查询之-AC自动机实现

本文章搬移改编自：blog.csdn.net/wyxeainn/ar…

#include <iostream>
#include <stdlib.h>
#include <queue>
#include <string>

using namespace std;

const int SON_NUM = 26;
std::string patten;  // 模式串
std::string text;    // 文本串
int match_res = 0;   // 匹配成功的数值

struct TrieNode
{  
    struct TrieNode* son[SON_NUM];  // 儿子们
    struct TrieNode* fail;          // 匹配失败时候的指针指向
    int num;                        // 以该节点所代表字符串为结尾的单词数
}* root;

// 创建节点
TrieNode* createNode()
{
    TrieNode* p;
    p = (TrieNode*)malloc(sizeof(TrieNode));
    for(int i = 0; i < SON_NUM; i++) p->son[i] = NULL;
    p->num = 0;
    p->fail = NULL;
    return p;
}

// 插入模式串，构建字典树
void insertPatten()
{
    TrieNode* p;
    p = root;
    int index = 0; 
    while(patten[index] != '\0')
    {
        int lowercase = patten[index] - 'a';
        if(p->son[lowercase] == NULL) {
            p->son[lowercase] = createNode();
        }
        p = p->son[lowercase];
        index++;
    }
    p->num++;
}

// 求fail指针。构造AC自动机。
void build_AC_automaton()
{
    TrieNode* p;
    p = root;
    queue<TrieNode*> qu;
    qu.push(p);
    while(!qu.empty()) {
        p = qu.front();
        qu.pop();
        for(int i = 0; i < SON_NUM; i++)
        {
            if(p->son[i] != NULL) // 第i个孩子存在
            {
                if(p == root) {
                    // p是根，根节点的孩子的失败指针都指向自己
                    p->son[i]->fail = root;
                } else {
                    TrieNode *node = p->fail;
                    while(node != NULL) {
                        if(node->son[i]!=NULL) {
                            p->son[i]->fail = node->son[i];
                            break;
                        }
                        node = node->fail;
                    }
                    if(node == NULL)
                        p->son[i]->fail = root;
                }
                qu.push(p->son[i]);
            }
        }
    }
}

void find_in_AC_automaton()
{
    TrieNode* p;
    p = root;
    int index = 0;
    while(text[index] != '\0')
    {
        int lowercase = text[index] - 'a';
        while(p->son[lowercase] == NULL && p != root)
            p = p->fail;   // 失配，转到能配的地方再尝试匹配

        p = p->son[lowercase];
        if(p == NULL) p = root;
        TrieNode* temp = p; // 把那些以当前节点的后缀作为后缀的字符串统计了
        while(temp != NULL && temp->num != -1)
        {
            match_res += temp->num;
            temp->num = -1;
            temp = temp->fail;
        }
        index++;
    }
}

// 记得释放接内存，用完及时归还系统，不然会爆的
void freeNode(TrieNode* node)
{
    if(node != NULL) {  
        for(int i = 0; i < SON_NUM; i++)
            freeNode(node->son[i]);
    }
    free(node);
}

int main()
{
    int n;
    cout << "input pattern num:";
    cin >> n;

    root = createNode();
    for(int i = 0; i < n; i++) {
        cout << "input the " << i + 1 << "th pattern:";
        cin >> patten;
        insertPatten();   // 用模式串构建字典树
    }

    cout << "input dst string:";
    cin >> text;

    build_AC_automaton(); // 构建AC自动机
    find_in_AC_automaton(); // 多模式匹配

    cout << "match num: " << match_res << endl;

    freeNode(root);        // 释放内存

    return 0;
}