Aho-Corasick(エイホ–コラシック法) (string/aho-corasick.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2024-04-27 01:27:28+09:00
• Include: #include "string/aho-corasick.hpp"

概要

複数文字列に対するパターンマッチングオートマトンを構築する.

Trie のそれぞれのノードがオートマトンのある状態に相当する.

使い方

• build(): 現在のTrieをもとにオートマトンを構築する.
• match(str, now): 現状態が now で, 文字列 str が現れた時に, 新たに各パターン文字列にマッチした回数を返す.
• move(c, now): 現状態が now で, 文字 c が現れたときに, 新たにパターン文字列にマッチした個数と, 次状態を返す.

計算量

• build(): パターン文字列の長さの総和
• クエリ: $O(1)$

Verified with

• test/verify/yukicoder-430.test.cpp

Code

/**
 * @brief Aho-Corasick(エイホ–コラシック法)
 *
 */
template <int char_size, int margin>
struct AhoCorasick : Trie<char_size + 1, margin> {
  using Trie<char_size + 1, margin>::Trie;

  const int FAIL = char_size;
  vector<int> correct;

  void build(bool heavy = true) {
    correct.resize(this->size());
    for (int i = 0; i < this->size(); i++) {
      correct[i] = (int)this->nodes[i].accept.size();
    }
    queue<int> que;
    for (int i = 0; i <= char_size; i++) {
      if (~this->nodes[0].nxt[i]) {
        this->nodes[this->nodes[0].nxt[i]].nxt[FAIL] = 0;
        que.emplace(this->nodes[0].nxt[i]);
      } else {
        this->nodes[0].nxt[i] = 0;
      }
    }
    while (!que.empty()) {
      auto &now = this->nodes[que.front()];
      int fail = now.nxt[FAIL];
      correct[que.front()] += correct[fail];
      que.pop();
      for (int i = 0; i < char_size; i++) {
        if (~now.nxt[i]) {
          this->nodes[now.nxt[i]].nxt[FAIL] = this->nodes[fail].nxt[i];
          if (heavy) {
            auto &u = this->nodes[now.nxt[i]].accept;
            auto &v = this->nodes[this->nodes[fail].nxt[i]].accept;
            vector<int> accept;
            set_union(begin(u), end(u), begin(v), end(v),
                      back_inserter(accept));
            u = accept;
          }
          que.emplace(now.nxt[i]);
        } else {
          now.nxt[i] = this->nodes[fail].nxt[i];
        }
      }
    }
  }

  map<int, int> match(const string &str, int now = 0) {
    map<int, int> result;
    for (auto &c : str) {
      now = this->nodes[now].nxt[c - margin];
      for (auto &v : this->nodes[now].accept) result[v] += 1;
    }
    return result;
  }

  pair<int64_t, int> move(const char &c, int now = 0) {
    now = this->nodes[now].nxt[c - margin];
    return {correct[now], now};
  }

  pair<int64_t, int> move(const string &str, int now = 0) {
    int64_t sum = 0;
    for (auto &c : str) {
      auto nxt = move(c, now);
      sum += nxt.first;
      now = nxt.second;
    }
    return {sum, now};
  }
};

#line 1 "string/aho-corasick.hpp"
/**
 * @brief Aho-Corasick(エイホ–コラシック法)
 *
 */
template <int char_size, int margin>
struct AhoCorasick : Trie<char_size + 1, margin> {
  using Trie<char_size + 1, margin>::Trie;

  const int FAIL = char_size;
  vector<int> correct;

  void build(bool heavy = true) {
    correct.resize(this->size());
    for (int i = 0; i < this->size(); i++) {
      correct[i] = (int)this->nodes[i].accept.size();
    }
    queue<int> que;
    for (int i = 0; i <= char_size; i++) {
      if (~this->nodes[0].nxt[i]) {
        this->nodes[this->nodes[0].nxt[i]].nxt[FAIL] = 0;
        que.emplace(this->nodes[0].nxt[i]);
      } else {
        this->nodes[0].nxt[i] = 0;
      }
    }
    while (!que.empty()) {
      auto &now = this->nodes[que.front()];
      int fail = now.nxt[FAIL];
      correct[que.front()] += correct[fail];
      que.pop();
      for (int i = 0; i < char_size; i++) {
        if (~now.nxt[i]) {
          this->nodes[now.nxt[i]].nxt[FAIL] = this->nodes[fail].nxt[i];
          if (heavy) {
            auto &u = this->nodes[now.nxt[i]].accept;
            auto &v = this->nodes[this->nodes[fail].nxt[i]].accept;
            vector<int> accept;
            set_union(begin(u), end(u), begin(v), end(v),
                      back_inserter(accept));
            u = accept;
          }
          que.emplace(now.nxt[i]);
        } else {
          now.nxt[i] = this->nodes[fail].nxt[i];
        }
      }
    }
  }

  map<int, int> match(const string &str, int now = 0) {
    map<int, int> result;
    for (auto &c : str) {
      now = this->nodes[now].nxt[c - margin];
      for (auto &v : this->nodes[now].accept) result[v] += 1;
    }
    return result;
  }

  pair<int64_t, int> move(const char &c, int now = 0) {
    now = this->nodes[now].nxt[c - margin];
    return {correct[now], now};
  }

  pair<int64_t, int> move(const string &str, int now = 0) {
    int64_t sum = 0;
    for (auto &c : str) {
      auto nxt = move(c, now);
      sum += nxt.first;
      now = nxt.second;
    }
    return {sum, now};
  }
};

Back to top page