首页 > 代码库 > 【AC自动机】AC自动机模板

【AC自动机】AC自动机模板

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <queue>
using namespace std;

const int N = 5e5+5, M = 1e6+5;
struct Trie {
    Trie *next[26];
    Trie *fail;
    int val;
}tree[N];
queue <Trie *> q;
int idx[128];

class ac_auto {
    private:
        int nxt;
        Trie *root;
    public:
        ac_auto() {
            nxt = 0;
            root = add();
        }
        Trie *add() {
            memset(&tree[nxt], 0, sizeof(Trie));
            tree[nxt].fail = root;
            return &tree[nxt++];
        }
        void insert(char *s) {
            Trie *rt = root;
            int len = strlen(s);
            for(int i = 0; i < len; i++) {
                int c = idx[s[i]+0];
                if(!rt->next[c])
                    rt->next[c] = add();
                rt = rt->next[c];
            }
            rt->val++;
        }
        void get_f() {
            queue <Trie *> q;
            q.push(root);
            root->fail = NULL;
            while(!q.empty()) {
                Trie *u = q.front(); q.pop();
                for(int c = 0; c < 26; c++) {
                    if(u->next[c]) {
                        Trie *f = u->fail;
                        while(f) {
                            if(f->next[c]) {
                                u->next[c]->fail = f->next[c];
                                break;
                            }
                            f = f->fail;
                        }
                        if(!f) u->next[c]->fail = root;
                        q.push(u->next[c]);
                    }
                }
            }
        }
        int match(char *s) {
            Trie *rt = root;
            int len = strlen(s), ret = 0;
            for(int i = 0; i < len; i++) {
                int c = idx[s[i]+0];
                while(!rt->next[c] && rt != root) rt = rt->fail;
                rt = rt->next[c];
                if(!rt) rt = root;
                Trie *p = rt;
                while(p != root) {
                    if(p->val) {
                        ret += p->val;
                        p->val = 0;
                    }
                    else break;
                    p = p->fail;
                }
            }
            return ret;
        }
};

void haxi()
{
    for(int i = 0; i < 26; i++) 
        idx['a'+i] = i;
}

int main()
{
    haxi();
    int T, m;
    char str[M];
    scanf("%d", &T);
    while(T--) {
        ac_auto ac;
        scanf("%d", &m);
        while(m--) {
            scanf("%s", str);
            ac.insert(str);
        }
        ac.get_f();
        scanf("%s", str);
        int ans = ac.match(str);
        printf("%d\n", ans);
    }
    return 0;
}

【AC自动机】AC自动机模板