競技プログラミング用 C++ ライブラリ
string/aho_corasick.hpp · GitHub で見る
共通ページ
複数のパターンを登録し、本文を1回走査して一致を調べます。重なる出現も数えます。
#include "lib/string/aho_corasick.hpp"
#include <cassert>
#include <string_view>
int main() {
AhoCorasick<> ac;
int he = ac.add("he");
int she = ac.add("she");
int e = ac.add("e");
ac.build();
std::string_view text = "ushe";
auto count = ac.occurrences(text);
assert(count[he] == 1 && count[she] == 1 && count[e] == 1);
assert(ac.count_matches(text) == 3);
ac.find_all(text, [&](std::size_t end, int id) {
std::size_t begin = end - ac.length(id);
assert(text.substr(begin, end - begin) == (id == he ? "he" : id == she ? "she" : "e"));
});
int state = ac.walk(ac.root, "us");
state = ac.transition(state, 'h');
state = ac.transition(state, 'e');
assert(ac.match_count(state) == 3);
}
add の返り値は追加順のパターン ID です。全てのパターンを追加してから build() を呼んでください。同じ文字列を複数回追加すると別の ID になり、それぞれを一致として数えます。build() 後にパターンを追加する場合は、clear() して全て登録し直します。本文ごとの検索は状態を変更しないので、同じオートマトンを繰り返し使えます。
| 操作 | 内容・返り値 | 計算量 |
|---|---|---|
AhoCorasick<Alphabet, First>() |
空の Trie。既定は英小文字26文字、First='a'
|
\(O(\sigma)\) |
reserve(total_length, patterns=0) |
パターンの長さの合計と個数を目安に領域を確保。呼び出しは任意 | \(O(V\sigma+P)\) |
clear() |
パターン・構築結果を全て削除。確保済み領域は保持 | \(O(V+P)\) |
add(pattern) |
string_view のパターンを追加し、パターン ID を返す。構築前のみ |
償却 \(O(\lvert pattern\rvert+K\sigma)\) |
build() |
suffix link・output link・全遷移を構築。構築済みなら何もしない | 初回 \(O(V\sigma)\)、以降 \(O(1)\) |
size() / empty()
|
パターンの個数/パターンが未登録か | \(O(1)\) |
node_count() / built()
|
根を含む Trie の頂点数/構築済みか | \(O(1)\) |
terminal(id) / length(id)
|
パターン ID の終端頂点/長さ | \(O(1)\) |
parent(v) / depth(v)
|
Trie の親/根からの深さ。根の親は根 | \(O(1)\) |
child(v,c) |
Trie の直接の子。存在しない、または文字が範囲外なら -1。構築前後で利用可能 | \(O(1)\) |
suffix_link(v) |
頂点の文字列の最長の真の接尾辞で、Trie に存在するものの頂点。根からは根 | \(O(1)\) |
output_link(v) |
suffix link をたどって最初に出会うパターンの終端頂点。自身は含まず、存在しなければ -1 | \(O(1)\) |
terminal_count(v) |
頂点 v ちょうどを終端とするパターン数。重複も含む |
\(O(1)\) |
match_count(v) |
頂点 v の文字列の末尾で一致するパターン数。自身・重複・空文字列も含む |
\(O(1)\) |
transition(v,c) |
状態 v で1文字読み込んだ次の状態。文字が範囲外なら根へ戻る |
\(O(1)\) |
walk(v,text) |
状態 v から本文を読み進め、最後の状態を返す |
\(O(\lvert text\rvert)\) |
for_each_match(v,visit) |
状態 v で一致する各パターン ID に visit(id) を呼ぶ |
\(O(Z_v)\) |
find_all(text,visit) |
全一致について visit(end,id) を呼ぶ。end は size_t の右端を含まない位置 |
\(O(n+Z)\) |
count_matches(text) |
全パターンの出現回数の合計を uint64_t で返す。列挙しない |
\(O(n)\) |
occurrences(text) |
各パターン ID の出現回数を vector<uint64_t> で返す。列挙しない |
\(O(n+V+P)\) |
bfs_order() |
根からの幅優先順を const vector<int>& で返す。根を含み、suffix link の行き先が先に来る |
\(O(1)\) |
propagate(values) |
頂点ごとの値を suffix link の行き先へ、幅優先順の逆順で += する |
\(O(V)\) |
root |
根の頂点番号 0 | \(O(1)\) |
表の \(K\) は、その add で新しく作られる頂点数です。パターンを追加した順に、まだ存在しない接頭辞へ頂点番号を割り当てます。terminal・length の引数はパターン ID、それ以外の状態を取る関数の引数は頂点番号です。
suffix_link・output_link・match_count と検索・集計・bfs_order・propagate は build() 後に使います。状態は読み込んだ文字列の接尾辞のうち、Trie に存在する最長のものを表します。本文を分割して読む場合は walk・transition の返り値を次の呼び出しへ渡します。
find_all の一致区間は [end-length(id),end) です。右端の昇順に列挙し、同じ右端では長いパターンから、同じパターンの重複は追加順の逆順に列挙します。コールバック中にオートマトンを変更しないでください。for_each_match も同じ長さ順です。
空文字列を追加すると、長さ \(n\) の本文の前後を含む \(n+1\) か所で一致します。空のパターン集合も使えます。総出現回数は uint64_t の範囲に収めてください。
propagate に渡す vector<T> の要素数は node_count() とし、T は += に対応させます。各走査状態に1を加えてから伝播すると、values[terminal(id)] が出現回数になります。空パターンも数えるときは、走査開始前に根にも1を加えます。入力配列を書き換えるため、同じ配列へ繰り返すと再加算されます。
文字はバイト単位です。add は [First,First+Alphabet) の文字だけを受け付けます。本文に範囲外の文字が出た場合は一致を途切れさせ、次の文字から再開します。AhoCorasick<256,0> なら NUL や128以上も含む全256バイト、AhoCorasick<10,'0'> なら数字を扱えます。パターン数・頂点数は INT_MAX 以下としてください。
遷移は配列で完成させるため、検索中は suffix link をたどらず1回の配列参照で移動します。一致 ID のリストは suffix link 間でコピーせず、output link を使います。再帰は使いません。既定の26文字では、頂点の配列と幅優先順だけで1頂点あたり約132バイトを使います。大きな文字集合を指定すると、その分メモリが増えます。
パターンの長さの合計を \(L\)、個数を \(P\)、Trie の頂点数を \(V\leq L+1\)、文字集合の大きさを \(\sigma\)、本文の長さを \(n\) とします。登録・構築は合計 \(O(L+P+V\sigma)\) 時間、保持する空間は \(O(V\sigma+P)\) です。occurrences は \(O(n+V+P)\) 時間・\(O(V+P)\) 追加空間、count_matches は \(O(n)\) 時間・\(O(1)\) 追加空間、全列挙は一致数を \(Z\) として \(O(n+Z)\) 時間・\(O(1)\) 追加空間です。コールバックと += の計算量は1回あたり \(O(1)\) としています。
#ifndef AHO_CORASICK_HPP
#define AHO_CORASICK_HPP
#include <array>
#include <bit>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <string_view>
#include <vector>
#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__))
#include <immintrin.h>
namespace aho_corasick_detail {
template<int Alphabet>
__attribute__((target("avx2"))) inline std::uint32_t complete(int *edges, const int *fallback) {
std::uint32_t mask = 0;
int i = 0;
for (; i + 8 <= Alphabet; i += 8) {
auto children = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(edges + i));
auto inherited = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(fallback + i));
auto missing = _mm256_cmpeq_epi32(children, _mm256_setzero_si256());
mask |= (std::uint32_t(~_mm256_movemask_ps(_mm256_castsi256_ps(missing))) & 255) << i;
_mm256_storeu_si256(reinterpret_cast<__m256i *>(edges + i), _mm256_blendv_epi8(children, inherited, missing));
}
for (; i < Alphabet; ++i) {
if (edges[i]) mask |= std::uint32_t(1) << i;
else edges[i] = fallback[i];
}
return mask;
}
}
#endif
template <int Alphabet = 26, int First = 'a'>
class AhoCorasick {
static_assert(0 < Alphabet && 0 <= First && First + Alphabet <= 256);
struct alignas(Alphabet >= 16 ? 64 : alignof(int)) Node {
std::array<int, Alphabet> next{};
int matches = 0;
int parent = 0, depth = 0, link = 0, output = -1, head = -1;
};
std::vector<Node> nodes_;
std::vector<int> terminal_, next_pattern_, order_;
bool built_ = false;
static unsigned symbol_(char c) { return static_cast<unsigned>(static_cast<unsigned char>(c) - First); }
void check_([[maybe_unused]] int v) const { assert(0 <= v && v < node_count()); }
void link_(int u, int f) {
nodes_[u].link = f;
nodes_[u].output = nodes_[f].head != -1 ? f : nodes_[f].output;
nodes_[u].matches += nodes_[f].matches;
order_.push_back(u);
}
public:
static constexpr int root = 0;
AhoCorasick() { nodes_.emplace_back(); }
void clear() {
nodes_.clear(), nodes_.emplace_back();
terminal_.clear(), next_pattern_.clear(), order_.clear();
built_ = false;
}
void reserve(int total_length, int patterns = 0) {
assert(0 <= total_length && total_length < std::numeric_limits<int>::max() && 0 <= patterns);
nodes_.reserve(std::size_t(total_length) + 1);
terminal_.reserve(patterns), next_pattern_.reserve(patterns);
}
int add(std::string_view pattern) {
assert(!built_ && size() < std::numeric_limits<int>::max());
int v = root;
for (char c : pattern) {
unsigned x = symbol_(c);
assert(x < Alphabet);
int u = nodes_[v].next[x];
if (!u) {
assert(node_count() < std::numeric_limits<int>::max());
u = node_count();
nodes_[v].next[x] = u;
nodes_.emplace_back();
nodes_[u].parent = v, nodes_[u].depth = nodes_[v].depth + 1;
}
v = u;
}
int id = size();
terminal_.push_back(v), next_pattern_.push_back(nodes_[v].head);
nodes_[v].head = id, ++nodes_[v].matches;
return id;
}
void build() {
if (built_) return;
order_.reserve(nodes_.size()), order_.push_back(root);
for (std::size_t i = 0; i < order_.size(); ++i) {
int v = order_[i];
auto &node = nodes_[v];
const auto &fallback = nodes_[node.link].next;
#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__))
if constexpr (8 <= Alphabet && Alphabet <= 32) {
if (__builtin_cpu_supports("avx2")) {
auto mask = aho_corasick_detail::complete<Alphabet>(node.next.data(), fallback.data());
while (mask) {
int c = std::countr_zero(mask), u = node.next[c], f = v == root ? root : fallback[c];
mask &= mask - 1;
link_(u, f);
}
continue;
}
}
#endif
for (int c = 0; c < Alphabet; ++c) {
if (int u = node.next[c]) {
link_(u, v == root ? root : fallback[c]);
} else {
node.next[c] = fallback[c];
}
}
}
built_ = true;
}
int size() const { return static_cast<int>(terminal_.size()); }
bool empty() const { return terminal_.empty(); }
int node_count() const { return static_cast<int>(nodes_.size()); }
bool built() const { return built_; }
int terminal(int id) const {
assert(0 <= id && id < size());
return terminal_[id];
}
int length(int id) const { return nodes_[terminal(id)].depth; }
int parent(int v) const { check_(v); return nodes_[v].parent; }
int depth(int v) const { check_(v); return nodes_[v].depth; }
int child(int v, char c) const {
check_(v);
unsigned x = symbol_(c);
if (x >= Alphabet) return -1;
int u = nodes_[v].next[x];
return u && nodes_[u].parent == v ? u : -1;
}
int suffix_link(int v) const { check_(v); assert(built_); return nodes_[v].link; }
int output_link(int v) const { check_(v); assert(built_); return nodes_[v].output; }
int terminal_count(int v) const {
check_(v);
return built_ && v ? nodes_[v].matches - nodes_[nodes_[v].link].matches : nodes_[v].matches;
}
int match_count(int v) const { check_(v); assert(built_); return nodes_[v].matches; }
int transition(int v, char c) const {
check_(v), assert(built_);
unsigned x = symbol_(c);
return x < Alphabet ? nodes_[v].next[x] : root;
}
int walk(int v, std::string_view text) const {
check_(v), assert(built_);
for (char c : text) v = transition(v, c);
return v;
}
const std::vector<int> &bfs_order() const { assert(built_); return order_; }
template <class F>
void for_each_match(int v, F &&visit) const {
check_(v), assert(built_);
if (nodes_[v].head == -1) v = nodes_[v].output;
for (; v != -1; v = nodes_[v].output)
for (int id = nodes_[v].head; id != -1; id = next_pattern_[id]) visit(id);
}
template <class F>
void find_all(std::string_view text, F &&visit) const {
assert(built_);
int v = root;
for_each_match(v, [&](int id) { visit(std::size_t(0), id); });
for (std::size_t i = 0; i < text.size(); ++i) {
v = transition(v, text[i]);
for_each_match(v, [&](int id) { visit(i + 1, id); });
}
}
std::uint64_t count_matches(std::string_view text) const {
assert(built_);
int v = root;
std::uint64_t count = nodes_[root].matches;
for (char c : text) {
v = transition(v, c);
count += nodes_[v].matches;
}
return count;
}
template <class T>
void propagate(std::vector<T> &values) const {
assert(built_ && values.size() == nodes_.size());
for (std::size_t i = order_.size(); i-- > 1;) {
int v = order_[i];
values[nodes_[v].link] += values[v];
}
}
std::vector<std::uint64_t> occurrences(std::string_view text) const {
assert(built_);
std::vector<std::uint64_t> hits(nodes_.size());
int v = root;
++hits[root];
for (char c : text) v = transition(v, c), ++hits[v];
propagate(hits);
std::vector<std::uint64_t> count(size());
for (int id = 0; id < size(); ++id) count[id] = hits[terminal_[id]];
return count;
}
};
#endif
#line 1 "string/aho_corasick.hpp"
#include <array>
#include <bit>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <string_view>
#include <vector>
#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__))
#include <immintrin.h>
namespace aho_corasick_detail {
template<int Alphabet>
__attribute__((target("avx2"))) inline std::uint32_t complete(int *edges, const int *fallback) {
std::uint32_t mask = 0;
int i = 0;
for (; i + 8 <= Alphabet; i += 8) {
auto children = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(edges + i));
auto inherited = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(fallback + i));
auto missing = _mm256_cmpeq_epi32(children, _mm256_setzero_si256());
mask |= (std::uint32_t(~_mm256_movemask_ps(_mm256_castsi256_ps(missing))) & 255) << i;
_mm256_storeu_si256(reinterpret_cast<__m256i *>(edges + i), _mm256_blendv_epi8(children, inherited, missing));
}
for (; i < Alphabet; ++i) {
if (edges[i]) mask |= std::uint32_t(1) << i;
else edges[i] = fallback[i];
}
return mask;
}
}
#endif
template <int Alphabet = 26, int First = 'a'>
class AhoCorasick {
static_assert(0 < Alphabet && 0 <= First && First + Alphabet <= 256);
struct alignas(Alphabet >= 16 ? 64 : alignof(int)) Node {
std::array<int, Alphabet> next{};
int matches = 0;
int parent = 0, depth = 0, link = 0, output = -1, head = -1;
};
std::vector<Node> nodes_;
std::vector<int> terminal_, next_pattern_, order_;
bool built_ = false;
static unsigned symbol_(char c) { return static_cast<unsigned>(static_cast<unsigned char>(c) - First); }
void check_([[maybe_unused]] int v) const { assert(0 <= v && v < node_count()); }
void link_(int u, int f) {
nodes_[u].link = f;
nodes_[u].output = nodes_[f].head != -1 ? f : nodes_[f].output;
nodes_[u].matches += nodes_[f].matches;
order_.push_back(u);
}
public:
static constexpr int root = 0;
AhoCorasick() { nodes_.emplace_back(); }
void clear() {
nodes_.clear(), nodes_.emplace_back();
terminal_.clear(), next_pattern_.clear(), order_.clear();
built_ = false;
}
void reserve(int total_length, int patterns = 0) {
assert(0 <= total_length && total_length < std::numeric_limits<int>::max() && 0 <= patterns);
nodes_.reserve(std::size_t(total_length) + 1);
terminal_.reserve(patterns), next_pattern_.reserve(patterns);
}
int add(std::string_view pattern) {
assert(!built_ && size() < std::numeric_limits<int>::max());
int v = root;
for (char c : pattern) {
unsigned x = symbol_(c);
assert(x < Alphabet);
int u = nodes_[v].next[x];
if (!u) {
assert(node_count() < std::numeric_limits<int>::max());
u = node_count();
nodes_[v].next[x] = u;
nodes_.emplace_back();
nodes_[u].parent = v, nodes_[u].depth = nodes_[v].depth + 1;
}
v = u;
}
int id = size();
terminal_.push_back(v), next_pattern_.push_back(nodes_[v].head);
nodes_[v].head = id, ++nodes_[v].matches;
return id;
}
void build() {
if (built_) return;
order_.reserve(nodes_.size()), order_.push_back(root);
for (std::size_t i = 0; i < order_.size(); ++i) {
int v = order_[i];
auto &node = nodes_[v];
const auto &fallback = nodes_[node.link].next;
#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__))
if constexpr (8 <= Alphabet && Alphabet <= 32) {
if (__builtin_cpu_supports("avx2")) {
auto mask = aho_corasick_detail::complete<Alphabet>(node.next.data(), fallback.data());
while (mask) {
int c = std::countr_zero(mask), u = node.next[c], f = v == root ? root : fallback[c];
mask &= mask - 1;
link_(u, f);
}
continue;
}
}
#endif
for (int c = 0; c < Alphabet; ++c) {
if (int u = node.next[c]) {
link_(u, v == root ? root : fallback[c]);
} else {
node.next[c] = fallback[c];
}
}
}
built_ = true;
}
int size() const { return static_cast<int>(terminal_.size()); }
bool empty() const { return terminal_.empty(); }
int node_count() const { return static_cast<int>(nodes_.size()); }
bool built() const { return built_; }
int terminal(int id) const {
assert(0 <= id && id < size());
return terminal_[id];
}
int length(int id) const { return nodes_[terminal(id)].depth; }
int parent(int v) const { check_(v); return nodes_[v].parent; }
int depth(int v) const { check_(v); return nodes_[v].depth; }
int child(int v, char c) const {
check_(v);
unsigned x = symbol_(c);
if (x >= Alphabet) return -1;
int u = nodes_[v].next[x];
return u && nodes_[u].parent == v ? u : -1;
}
int suffix_link(int v) const { check_(v); assert(built_); return nodes_[v].link; }
int output_link(int v) const { check_(v); assert(built_); return nodes_[v].output; }
int terminal_count(int v) const {
check_(v);
return built_ && v ? nodes_[v].matches - nodes_[nodes_[v].link].matches : nodes_[v].matches;
}
int match_count(int v) const { check_(v); assert(built_); return nodes_[v].matches; }
int transition(int v, char c) const {
check_(v), assert(built_);
unsigned x = symbol_(c);
return x < Alphabet ? nodes_[v].next[x] : root;
}
int walk(int v, std::string_view text) const {
check_(v), assert(built_);
for (char c : text) v = transition(v, c);
return v;
}
const std::vector<int> &bfs_order() const { assert(built_); return order_; }
template <class F>
void for_each_match(int v, F &&visit) const {
check_(v), assert(built_);
if (nodes_[v].head == -1) v = nodes_[v].output;
for (; v != -1; v = nodes_[v].output)
for (int id = nodes_[v].head; id != -1; id = next_pattern_[id]) visit(id);
}
template <class F>
void find_all(std::string_view text, F &&visit) const {
assert(built_);
int v = root;
for_each_match(v, [&](int id) { visit(std::size_t(0), id); });
for (std::size_t i = 0; i < text.size(); ++i) {
v = transition(v, text[i]);
for_each_match(v, [&](int id) { visit(i + 1, id); });
}
}
std::uint64_t count_matches(std::string_view text) const {
assert(built_);
int v = root;
std::uint64_t count = nodes_[root].matches;
for (char c : text) {
v = transition(v, c);
count += nodes_[v].matches;
}
return count;
}
template <class T>
void propagate(std::vector<T> &values) const {
assert(built_ && values.size() == nodes_.size());
for (std::size_t i = order_.size(); i-- > 1;) {
int v = order_[i];
values[nodes_[v].link] += values[v];
}
}
std::vector<std::uint64_t> occurrences(std::string_view text) const {
assert(built_);
std::vector<std::uint64_t> hits(nodes_.size());
int v = root;
++hits[root];
for (char c : text) v = transition(v, c), ++hits[v];
propagate(hits);
std::vector<std::uint64_t> count(size());
for (int id = 0; id < size(); ++id) count[id] = hits[terminal_[id]];
return count;
}
};