Aho-Corasick
(String/ahocorasick.hpp)

• View this file on GitHub
• Last update: 2024-01-12 04:16:01+09:00
• Include: #include "String/ahocorasick.hpp"

Depends on

• Trie (String/trie.hpp)

Code

#pragma once
#include "String/trie.hpp"

template<int SZ,int start>struct AhoCorasick:Trie<SZ,start>{
    using super=Trie<SZ,start>;
    using super::vs;
    vector<int> fail,bfs;
    void make_failure(){
        fail.resize(vs.size());
        queue<int> que;
        que.push(0);
        while(!que.empty()){
            int v=que.front();
            que.pop();
            bfs.push_back(v);
            rep(c,0,SZ){
                if(vs[v].nxt[c]==-1)continue;
                int to=vs[v].nxt[c];
                que.push(to);
                if(v==0)continue;
                int f=fail[v];
                while(f and vs[f].nxt[c]==-1)f=fail[f];
                if(vs[f].nxt[c]!=-1)fail[to]=vs[f].nxt[c];
            }
        }
    }
    void make_next(){
        for(auto& v:bfs){
            rep(c,0,SZ)if(vs[v].nxt[c]==-1){
                vs[v].nxt[c]=vs[fail[v]].nxt[c];
                if(vs[v].nxt[c]==-1)vs[v].nxt[c]=0;
            }
        }
    }
    void merge_id(){
        for(auto& v:bfs)if(v>0){
            rep(c,0,SZ){
                int from=vs[fail[v]].nxt[c];
                int to=vs[v].nxt[c];
                if(from!=-1 and to!=-1){
                    for(auto& id:vs[from].ids)vs[to].ids.push_back(id);
                }
            }
        }
    }
};


/**
 * @brief Aho-Corasick
*/

#line 2 "String/trie.hpp"

template<int SZ,int start>struct Trie{
    struct Node{
        array<int,SZ> nxt;
        vector<int> ids;
        Node(){fill(ALL(nxt),-1);}
    };
    vector<Node> vs;
    
    Trie(){vs.push_back(Node());}
    Node operator[](const int& i)const{return vs[i];}
    int move(int pos,char c){
        return pos<0?-1:vs[pos].nxt[int(c-start)];
    }
    void add(const string& s,int id){
        int pos=0;
        for(auto& c:s){
            int to=move(pos,c);
            if(to==-1){
                to=vs.size();
                vs.push_back(Node());
                vs[pos].nxt[int(c-start)]=to;
            }
            pos=to;
        }
        vs[pos].ids.push_back(id);
    }
    int find(string& s){
        int pos=0;
        for(auto& c:s){
            pos=move(pos,c);
            if(pos==-1)return -1;
        }
        return pos;
    }
};

/**
 * @brief Trie
*/
#line 3 "String/ahocorasick.hpp"

template<int SZ,int start>struct AhoCorasick:Trie<SZ,start>{
    using super=Trie<SZ,start>;
    using super::vs;
    vector<int> fail,bfs;
    void make_failure(){
        fail.resize(vs.size());
        queue<int> que;
        que.push(0);
        while(!que.empty()){
            int v=que.front();
            que.pop();
            bfs.push_back(v);
            rep(c,0,SZ){
                if(vs[v].nxt[c]==-1)continue;
                int to=vs[v].nxt[c];
                que.push(to);
                if(v==0)continue;
                int f=fail[v];
                while(f and vs[f].nxt[c]==-1)f=fail[f];
                if(vs[f].nxt[c]!=-1)fail[to]=vs[f].nxt[c];
            }
        }
    }
    void make_next(){
        for(auto& v:bfs){
            rep(c,0,SZ)if(vs[v].nxt[c]==-1){
                vs[v].nxt[c]=vs[fail[v]].nxt[c];
                if(vs[v].nxt[c]==-1)vs[v].nxt[c]=0;
            }
        }
    }
    void merge_id(){
        for(auto& v:bfs)if(v>0){
            rep(c,0,SZ){
                int from=vs[fail[v]].nxt[c];
                int to=vs[v].nxt[c];
                if(from!=-1 and to!=-1){
                    for(auto& id:vs[from].ids)vs[to].ids.push_back(id);
                }
            }
        }
    }
};


/**
 * @brief Aho-Corasick
*/

Back to top page