cp-library

競技プログラミング用 C++ ライブラリ

View the Project on GitHub qandori/cp-library

:heavy_check_mark: Aho–Corasick

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
Back to top page