src/string/aho.hpp

• View this file on GitHub
• Last update: 2025-12-30 15:03:15+09:00
• Include: #include "src/string/aho.hpp"

Depends on

• src/template.hpp

Code

#pragma once
#include "../template.hpp"

struct Node {
    int prefix_cnt;
    int parent;
    vector<int> nxt,ids;
    Node(int sigma=26,int parent=-1):prefix_cnt(0),parent(-1),nxt(sigma,-1){}
};

struct Trie {
    int sigma;
    vector<Node> nodes;

    Trie(int _sigma):sigma(_sigma),nodes(1,Node(sigma)){}
    int add(vector<int> V){
        int id=nodes[0].prefix_cnt;
        int now=0;
        for(int v: V){
            if(nodes[now].nxt[v]==-1){
                nodes[now].nxt[v]=ssize(nodes);
                nodes.push_back(Node(sigma,now));
            }
            nodes[now].prefix_cnt++;
            now=nodes[now].nxt[v];
        }
        nodes[now].ids.push_back(id);
        return now;
    }
};

struct AhoCorasick:Trie {
    vector<int> cnt;
    AhoCorasick(int _sigma):Trie(_sigma+1){}
    void build(){
        cnt.assign(ssize(nodes),0);
        rep(i,ssize(nodes)) cnt[i]=ssize(nodes[i].ids);
        queue<int> q;
        rep(i,sigma){
            if(~nodes[0].nxt[i]){
                nodes[nodes[0].nxt[i]].nxt[sigma]=0;
                q.push(nodes[0].nxt[i]);
            }
            else{
                nodes[0].nxt[i]=0;
            }
        }
        while(!q.empty()){
            auto& now=nodes[q.front()];
            int fail=now.nxt[sigma];
            cnt[q.front()]|=cnt[fail];
            q.pop();
            rep(i,sigma){
                if(~now.nxt[i]){
                    nodes[now.nxt[i]].nxt[sigma]=nodes[fail].nxt[i];
                    q.push(now.nxt[i]);
                }
                else{
                    now.nxt[i]=nodes[fail].nxt[i];
                }
            }
        }
    }

    //[cnt, next]
    pair<int,int> move(int c,int now){
        now=nodes[now].nxt[c];
        return {cnt[now],now};
    }

    int size(){
        return ssize(nodes);
    }
};

#line 2 "src/template.hpp"
#include <bits/stdc++.h>
using namespace std;
#define si(a) (long)a.size()
#define fi first
#define se second
#define all(x) x.begin(),x.end()
#define rep(i,n) for(int i=0;i<(int)(n);++i)
template<typename S,typename F> bool chmin(S&a,F b){return b<a?(a=b,1):0;}
template<typename S,typename F> bool chmax(S&a,F b){return b>a?(a=b,1):0;}
bool _=(ios::sync_with_stdio(0),cin.tie(0),cout<<fixed<<setprecision(16),0);
#line 3 "src/string/aho.hpp"

struct Node {
    int prefix_cnt;
    int parent;
    vector<int> nxt,ids;
    Node(int sigma=26,int parent=-1):prefix_cnt(0),parent(-1),nxt(sigma,-1){}
};

struct Trie {
    int sigma;
    vector<Node> nodes;

    Trie(int _sigma):sigma(_sigma),nodes(1,Node(sigma)){}
    int add(vector<int> V){
        int id=nodes[0].prefix_cnt;
        int now=0;
        for(int v: V){
            if(nodes[now].nxt[v]==-1){
                nodes[now].nxt[v]=ssize(nodes);
                nodes.push_back(Node(sigma,now));
            }
            nodes[now].prefix_cnt++;
            now=nodes[now].nxt[v];
        }
        nodes[now].ids.push_back(id);
        return now;
    }
};

struct AhoCorasick:Trie {
    vector<int> cnt;
    AhoCorasick(int _sigma):Trie(_sigma+1){}
    void build(){
        cnt.assign(ssize(nodes),0);
        rep(i,ssize(nodes)) cnt[i]=ssize(nodes[i].ids);
        queue<int> q;
        rep(i,sigma){
            if(~nodes[0].nxt[i]){
                nodes[nodes[0].nxt[i]].nxt[sigma]=0;
                q.push(nodes[0].nxt[i]);
            }
            else{
                nodes[0].nxt[i]=0;
            }
        }
        while(!q.empty()){
            auto& now=nodes[q.front()];
            int fail=now.nxt[sigma];
            cnt[q.front()]|=cnt[fail];
            q.pop();
            rep(i,sigma){
                if(~now.nxt[i]){
                    nodes[now.nxt[i]].nxt[sigma]=nodes[fail].nxt[i];
                    q.push(now.nxt[i]);
                }
                else{
                    now.nxt[i]=nodes[fail].nxt[i];
                }
            }
        }
    }

    //[cnt, next]
    pair<int,int> move(int c,int now){
        now=nodes[now].nxt[c];
        return {cnt[now],now};
    }

    int size(){
        return ssize(nodes);
    }
};

Back to top page