Aho-Corasickに動的にキーワードを追加する

この記事はIQが1Advent Calendarの1日目の記事になります.
「起床時間とか就寝時間が1日目の記事になる」と言っていましたが,ギリギリ実装が間に合った(バグってるかもしれないが)ので起床時間と就寝時間はやめて文字列照合アルゴリズムについて紹介することにしました.
IQ1でもわかるように書くつもりです(これは嘘でIQ1なのでまともな文章が書けない)
(AhoなのでIQ1と指摘を受けました. IQ1はAhoですがAhoさんはwikiを見てもわかるとおりすごい人です.IQが異常に高そうです)

Aho-Corasickについて

複数のキーワードに対する文字列照合アルゴリズムとしてAho-Corasickというものがある.
Aho-Corasickについては Alfred V.Aho and Margaret J.Corasick Bell Laboratories "Efficient String Matching: An Aid to Bibliographic Search"という論文にかかれている.
このキーワード群を静的に構築する方法は下のブログに書いてある.
d.hatena.ne.jp
これのC++の実装はalgoogle.
algoogle.hadrori.jp

簡単説明するとTrieを構築した後にTrieでの検索に失敗した際の遷移関数を定義することで効率的に文字列照合を行うというものである.

このブログではキーワードを動的に追加する方法を紹介する.

キーワードの動的な追加方法

共立出版からでている「情報検索アルゴリズム」を参考に...

言葉の定義

goto関数 : 検索に失敗しない場合の遷移
failure関数 : 検索に失敗した際の遷移

この関数については上のブログでも定義されてる...(そっち見たほうがいい)
Aho-Corasickはこのfailure関数を幅優先探索で構築する.

これを動的に行う場合,更新による影響が及ぶ部分だけを更新すればよい.
更新による影響が及ぶ範囲というのは,新しく追加された部分にfailure関数が到達する部分である(これによりoutput関数も更新されるので).
下の図ではノード9,10に相当する.

これを動的に行うには,2つの段階が必要.
1. 追加するキーワードに対応するノードを作成し,追加されたノードのfailure関数とoutput関数を更新する
2. 既存のグラフのfailure関数とoutput関数を更新する

例えば...
既に構築されているキーワード群は { "ab", "abcde" , "bc" , "bab" , "d"}とし,新しく "bcde"を追加するとします

1ステップで赤色の更新を, 2ステップで緑色の更新を行います(output関数の更新は省略しました)
f:id:chakku000:20171201003835p:plain

ステップ1

赤いノードの追加 : Trieと同様なので省略
赤いfailure関数の追加
f:id:chakku000:20171201004809p:plain
node1->node2はfailure関数を,赤い線は追加されるnode2->node4というfailure関数を意味する.
薄い水色の部分は同じ文字の並びを意味する.

上の図の関係を見つければ,赤い辺をfailure関数として定めることができる.(多分ここが１番重要.ステップ2もこの考え方)

ステップ2

緑のノードの追加.
ここでも2つめの図と同じような考察をする.
例えば1つめの図のノード8のfailure関数は4を指す(f(8) = 4) .
そうすると
8 <-> node1
4 <-> node3
9 <-> node2
11 <-> node4
の対応が見え,図1における 9 -> 11 (つまり図2におけるnode2->node4)を貼ることができる.

図1においてこれを行うためにはノード4がノード8のfailure遷移先であることをノード4は知っている必要がある.そのためfailure関数の逆関数,逆failure関数も保持しておく必要がある.

ステップ2は繰り返し適用することで更新が必要な状態は全て更新される.
この繰り返し更新するのはqueueを用いて行えばよい

実装

C++における実装. addが動的にキーワードを追加する部分.
それ以外ははどろりさんの実装を真似ている.

class AhoCorasick{
    private:
        const int nodeSize = 30;
    public:
        vector<AhoCorasick*> nodes;
        AhoCorasick* failure;
        set<AhoCorasick*> rev_failure;
        vector<int> matched;    // 対応ノードにpattern[i]がマッチした
        vector<string> pattern;
        AhoCorasick() : nodes(nodeSize),failure(nullptr){
            for(int i=0;i<nodeSize;i++) nodes[i] = nullptr;
        }
        void rootinit(){
            for(int i=0;i<nodeSize;i++) nodes[i] = this;
        }

        // 静的に追加(from algoogle)
        void insert(const vector<string> pattern_){
            pattern = pattern_;
            // Trie構築
            AhoCorasick *r = this;
            failure = this;
            rev_failure.insert(this);
            for(int i=0;i<pattern.size();i++){
                r = this;
                for(int j=0;j<pattern[i].size();j++){
                    int c = pattern[i][j]-'a';
                    if(!r->nodes[c]) r->nodes[c] = new AhoCorasick;
                    r = r->nodes[c];
                }
                r->matched.push_back(i);
            }

            // 辺生成
            queue<AhoCorasick*> que;
            for(int i=0;i<nodeSize;i++){
                if(!this->nodes[i]) this->nodes[i] = this; // 存在しないノードのfailure
                else{
                    this->nodes[i]->failure = this; // 深さ1のfailure関数は全てrootに接続される
                    this->rev_failure.insert(this->nodes[i]);
                    que.push(this->nodes[i]);
                }
            }

            // failure関数の更新とmatchedの更新
            while(!que.empty()){
                r = que.front();que.pop();
                for(int i=0;i<nodeSize;i++) if(r->nodes[i]){
                    AhoCorasick* sfx = r->failure;
                    while(!sfx->nodes[i]) sfx = sfx->failure;
                    r->nodes[i]->failure = sfx->nodes[i];
                    sfx->nodes[i]->rev_failure.insert(r->nodes[i]);
                    copy(sfx->nodes[i]->matched.begin(),sfx->nodes[i]->matched.end(),back_inserter(r->nodes[i]->matched));
                    que.push(r->nodes[i]);
                }
            }
        }

        // 動的に追加
        void add(const string& keyword){
            int idx=0;
            int depth=0;  // 遷移回数
            AhoCorasick *r = this;
            vector<AhoCorasick*> s(keyword.size()+1,nullptr);   // s[i] : keyword[i-1]で遷移した先
            s[0] = this;
            // rootからの遷移は深さ1以上の状態からの遷移とことなる特徴を持つので特別に扱う
            if(nodes[keyword[0]-'a'] == this || nodes[keyword[0]-'a'] == nullptr){  // 対応するgoto関数が定義されているか?
                // 先頭文字に対するgoto関数が定義されていない場合,深さ1の状態だけは生成
                nodes[keyword[0]-'a'] = new AhoCorasick;
                s[1] = nodes[keyword[0]-'a'];

                nodes[keyword[0]-'a']->failure = this;      // 深さ1のfailure関数は必ずroot
                this->rev_failure.insert(nodes[keyword[0]-'a']);

                idx = 1;
            }else{
                // 先頭文字に対する状態遷移がある場合はgoto関数がfailになるまで既存の状態遷移をたどる
                while(idx < keyword.size() and s[idx]->nodes[keyword[idx]-'a']){
                    s[idx+1] = s[idx]->nodes[keyword[idx]-'a'];
                    idx++;
                }
                depth = idx;
            }


            //状態遷移の追加とfailure関数,matchedの更新
            for(int i=idx;i<keyword.size();i++){
                // 新しいノード作成
                s[i]->nodes[keyword[i]-'a'] = new AhoCorasick;
                s[i+1] = s[i]->nodes[keyword[i]-'a'];

                // 新しいノードの親ノードのfailureを辿る
                AhoCorasick *state = s[i]->failure;
                while(!state->nodes[keyword[i]-'a']) state = state->failure;

                // 新しいノードのfailure関数を設定
                s[i+1]->failure = state->nodes[keyword[i]-'a'];
                state->nodes[keyword[i]-'a']->rev_failure.insert(s[i+1]);

                // 新しいノードのmatchedを決定
                s[i+1]->matched = state->nodes[keyword[i]-'a']->matched;
            }

            // 追加ノードが無い場合のs[n]を定義
            if(!s[keyword.size()]) s[keyword.size()] = s[keyword.size()-1]->nodes[keyword[keyword.size()-1]-'a'];
            s[keyword.size()]->matched.push_back(pattern.size());

            // 既存状態の更新

            // 追加状態の親はs[depth]
            r = s[depth];

            queue<pair<AhoCorasick*,int>> que;
            for(AhoCorasick* node : r->rev_failure){
                que.push(make_pair(node,depth));
            }

            while(!que.empty()){
                r = que.front().first;
                int i = que.front().second;
                que.pop();
                while(i<keyword.size() and r->nodes[keyword[i]-'a']){
                    AhoCorasick* state = r->nodes[keyword[i]-'a'];
                    state->failure->rev_failure.erase(state);
                    state->failure = s[i+1];
                    s[i+1]->rev_failure.insert(state);
                    r = state;
                    i++;
                }

                if(i>=keyword.size()) r->matched.push_back(pattern.size());

                for(AhoCorasick* node : r->rev_failure){
                    que.push(make_pair(node,i));
                }
            }
            pattern.push_back(keyword);
        }

        // ret[tのindex] := t[index]から開始する文字列はpattern[ret[index][i]]にマッチする
        // 例えば pattern[0] = a , pattern[1] = ab , t = ab なら ret[0] = {0(a),1(ab)}
        vector<vector<int>> match(const string& s){
            AhoCorasick* r = this;
            vector<vector<int>> ret(s.size());
            //map<int,vector<int>> ret;
            for(int i=0;i<s.size();i++){
                int c = s[i] - 'a';
                while(!r->nodes[c]){
                    r = r->failure;
                }
                r = r->nodes[c];
                for(int j=0;j<r->matched.size();j++){
                    int sz = pattern[r->matched[j]].size();
                    ret[i-sz+1].pb(r->matched[j]);
                }
            }
            return ret;
        }
};