cp-library

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

View the Project on GitHub qandori/cp-library

:heavy_check_mark: Suffix Automaton(接尾辞オートマトン)

string/suffix_automaton.hpp · GitHub で見る

共通ページ

末尾へ文字を追加しながら、部分文字列の検索、出現回数、相異なる部分文字列の数を求めます。NUL・128 以上の値を含む全 256 種類のバイトに対応します。

#include "lib/string/suffix_automaton.hpp"

int main() {
    SuffixAutomaton sam("ababa");
    bool found = sam.contains("bab"); // true
    auto count = sam.count("aba"); // 2
    auto distinct = sam.distinct_count(); // 9
    sam.push_back('b');
    count = sam.count("ab"); // 3
    auto match = sam.longest_common_substring("zzbaba");
    // match.text_pos = 1, query_pos = 2, length = 4
}
操作 内容・返り値
SuffixAutomaton() / SuffixAutomaton(s) 空の文字列/std::string_view から構築
clear() 空へ戻す。確保済みの領域は保持
reserve(n) 長さ n の文字列に必要な領域を予約
push_back(c) char を末尾に追加し、追加後の last() を int で返す
append(s) std::string_view を末尾へ追加。返り値は void
size() / empty() 追加したバイト数を int/空かどうかを bool で返す
distinct_count() 異なる空でない部分文字列の数を std::uint64_t で返す
find(s) / contains(s) 部分文字列の到達状態を int/存在するかを bool で返す。不在の状態番号は -1
count(s) 重なりを含む出現回数を std::uint64_t で返す。不在なら 0
longest_common_substring(s) 構築した文字列と s の最長共通部分文字列の位置・長さを Match で返す

Match は int text_pos, query_pos, length を持ちます。構築した文字列の [text_pos,text_pos+length) と、引数の [query_pos,query_pos+length) が一致します。最長の候補が複数ある場合は、引数での出現位置が最小のものを返します。その部分文字列の構築文字列での出現位置も最小です。一致がなければ全フィールドが 0 です。

空文字列は常に存在し、find("") は root、count("") は size()+1 です。distinct_count() には空文字列を含めません。

状態を使った DP や出現回数の伝播には次の API を使えます。

操作・定数 内容・返り値
root 初期状態の番号 0
node_count() 初期状態を含む状態数を int で返す
last() 現在の文字列全体の状態を int で返す。空では root
length(v) / min_length(v) 状態 v が表す部分文字列の最大長/最小長を int で返す
suffix_link(v) suffix link の状態番号を int で返す。初期状態では -1
transition(v,c) 1 バイトの遷移先を int で返す。存在しなければ -1
transitions(v) (バイト,遷移先) の std::vector<std::pair<unsigned char,int>> を返す。順序は未規定
walk(v,s) v から s を読んだ状態を int で返す。途中で遷移できなければ -1
first_occurrence(v) 状態の最大長の文字列が最初に現れる [l,r) を std::pair<int,int> で返す。初期状態では (0,0)
occurrences() 各状態の出現回数を std::vector<std::uint64_t> で返す。初期状態は size()+1
prefix_nodes() 各非空接頭辞の状態番号を const std::vector<int>& で返す。要素 i は長さ i+1 の接頭辞
order_by_length() length の昇順に状態番号を並べた std::vector<int> を返す

状態は同じ終了位置集合を持つ部分文字列をまとめたものです。状態 v の出現回数は、長さが min_length(v) 以上 length(v) 以下の、その状態に到達するどの部分文字列にも共通です。

初期状態以外の suffix link は最大長がより短い状態を指し、文字の遷移は最大長がより長い状態を指します。状態番号の大小は suffix link の向きと一致するとは限りません。order_by_length() は suffix link の伝播や遷移上の DP に使えます。

出現回数は必要になったときに一括計算し、文字追加で無効化します。繰り返し取得しても回数を重複加算せず、取得後も文字を追加できます。状態番号は追加しても変わりませんが、suffix link・遷移先・最小長・出現回数は変わることがあります。clear() では状態番号もリセットします。コピーは独立します。

文字数・予約長は (INT_MAX-1)/3 以下とします。状態を引数に取る操作には [0,node_count()) の番号を渡してください。最長共通部分文字列の引数は INT_MAX バイト以下です。削除や途中への挿入には対応しません。バイト列として処理するため、UTF-8 の文字単位ではありません。返した配列は値のコピーです。ただし prefix_nodes() の参照とその要素の有効期間には std::vector と同じ規則が適用されます。

文字数を \(N\)、検索する文字列の長さを \(M\)、登場するバイトの種類数を \(\sigma\le256\) とします。構築・文字追加の合計は \(O(N\sigma)\) 時間、保持空間は \(O(N)\) です。push_back は償却 \(O(\sigma)\)、transition は \(O(\sigma)\)、find・contains・walk・最長共通部分文字列は \(O(M\sigma)\) です。count は文字追加後の初回に追加で \(O(N)\)、以降は \(O(M\sigma)\) 時間です。occurrences・order_by_length は \(O(N)\) 時間・追加空間、transitions(v) は出次数に比例する時間・追加空間です。reserve は予約長と既存要素数に比例する時間、clear は \(O(N)\) 時間、その他の取得は \(O(1)\) 時間です。

検証コード

コード

#ifndef SUFFIX_AUTOMATON_HPP
#define SUFFIX_AUTOMATON_HPP

#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <limits>
#include <string_view>
#include <utility>
#include <vector>

class SuffixAutomaton {
	struct Node { int length, link, first_end, hits, head; };
	struct Edge { int to, next; unsigned char symbol; };
	static constexpr int maximum_size = (std::numeric_limits<int>::max() - 1) / 3;
	std::vector<Node> nodes_;
	std::vector<Edge> edges_;
	std::vector<int> prefix_;
	std::array<int, 256> root_edges_;
	mutable std::vector<std::uint64_t> counts_;
	std::uint64_t distinct_ = 0;
	int last_ = 0;
	void check_([[maybe_unused]] int v) const { assert(0 <= v && v < node_count()); }
	int edge_(int v, unsigned char c) const {
		if (!v) return root_edges_[c];
		for (int e = nodes_[v].head; e != -1; e = edges_[e].next)
			if (edges_[e].symbol == c) return e;
		return -1;
	}
	void add_edge_(int v, unsigned char c, int to) {
		int e = static_cast<int>(edges_.size());
		edges_.push_back({to, nodes_[v].head, c});
		nodes_[v].head = e;
		if (!v) root_edges_[c] = e;
	}
	void prepare_counts_() const {
		if (!counts_.empty()) return;
		auto order = order_by_length();
		counts_.resize(nodes_.size());
		for (int v = 0; v < node_count(); ++v) counts_[v] = nodes_[v].hits;
		for (auto it = order.rbegin(); it != order.rend(); ++it)
			if (*it) counts_[nodes_[*it].link] += counts_[*it];
		counts_[root] = std::uint64_t(size()) + 1;
	}

public:
	struct Match { int text_pos = 0, query_pos = 0, length = 0; };
	static constexpr int root = 0;
	SuffixAutomaton() { clear(); }
	explicit SuffixAutomaton(std::string_view s) : SuffixAutomaton() {
		assert(s.size() <= std::size_t(maximum_size));
		reserve(static_cast<int>(s.size()));
		append(s);
	}
	void clear() {
		nodes_.clear(), edges_.clear(), prefix_.clear(), counts_.clear();
		nodes_.push_back({0, -1, 0, 0, -1});
		root_edges_.fill(-1);
		last_ = root, distinct_ = 0;
	}
	void reserve(int n) {
		assert(0 <= n && n <= maximum_size);
		nodes_.reserve(std::size_t(n) * 2 + 1);
		edges_.reserve(std::size_t(n) * 3 + 1);
		prefix_.reserve(n);
	}
	int size() const { return static_cast<int>(prefix_.size()); }
	bool empty() const { return prefix_.empty(); }
	int node_count() const { return static_cast<int>(nodes_.size()); }
	int last() const { return last_; }
	std::uint64_t distinct_count() const { return distinct_; }
	int length(int v) const { check_(v); return nodes_[v].length; }
	int min_length(int v) const { check_(v); return v ? nodes_[nodes_[v].link].length + 1 : 0; }
	int suffix_link(int v) const { check_(v); return nodes_[v].link; }
	std::pair<int, int> first_occurrence(int v) const {
		check_(v);
		return {nodes_[v].first_end - nodes_[v].length, nodes_[v].first_end};
	}
	int transition(int v, char c) const {
		check_(v);
		int e = edge_(v, static_cast<unsigned char>(c));
		return e == -1 ? -1 : edges_[e].to;
	}
	std::vector<std::pair<unsigned char, int>> transitions(int v) const {
		check_(v);
		std::vector<std::pair<unsigned char, int>> result;
		for (int e = nodes_[v].head; e != -1; e = edges_[e].next) result.emplace_back(edges_[e].symbol, edges_[e].to);
		return result;
	}
	int push_back(char symbol) {
		assert(size() < maximum_size);
		counts_.clear();
		unsigned char c = static_cast<unsigned char>(symbol);
		int cur = node_count(), p = last_, e = -1;
		nodes_.push_back({nodes_[last_].length + 1, 0, size() + 1, 1, -1});
		while (p != -1 && (e = edge_(p, c)) == -1) {
			add_edge_(p, c, cur);
			p = nodes_[p].link;
		}
		if (p != -1) {
			int q = edges_[e].to;
			if (nodes_[p].length + 1 == nodes_[q].length) nodes_[cur].link = q;
			else {
				int clone = node_count();
				nodes_.push_back({nodes_[p].length + 1, nodes_[q].link, nodes_[q].first_end, 0, -1});
				for (int at = nodes_[q].head; at != -1; at = edges_[at].next)
					add_edge_(clone, edges_[at].symbol, edges_[at].to);
				while (p != -1 && (e = edge_(p, c)) != -1 && edges_[e].to == q) {
					edges_[e].to = clone;
					p = nodes_[p].link;
				}
				nodes_[q].link = nodes_[cur].link = clone;
			}
		}
		last_ = cur;
		prefix_.push_back(cur);
		distinct_ += nodes_[cur].length - nodes_[nodes_[cur].link].length;
		return cur;
	}
	void append(std::string_view s) {
		assert(s.size() <= std::size_t(maximum_size - size()));
		for (char c : s) push_back(c);
	}
	int walk(int v, std::string_view s) const {
		check_(v);
		for (char c : s) { v = transition(v, c); if (v == -1) break; }
		return v;
	}
	int find(std::string_view s) const { return walk(root, s); }
	bool contains(std::string_view s) const { return find(s) != -1; }
	std::uint64_t count(std::string_view s) const {
		int v = find(s);
		if (v == -1) return 0;
		prepare_counts_();
		return counts_[v];
	}
	std::vector<std::uint64_t> occurrences() const { prepare_counts_(); return counts_; }
	const std::vector<int> &prefix_nodes() const { return prefix_; }
	std::vector<int> order_by_length() const {
		std::vector<int> buckets(std::size_t(size()) + 1), order(nodes_.size());
		for (const auto &node : nodes_) ++buckets[node.length];
		for (int i = 1; i <= size(); ++i) buckets[i] += buckets[i - 1];
		for (int v = node_count() - 1; v >= 0; --v) order[--buckets[nodes_[v].length]] = v;
		return order;
	}
	Match longest_common_substring(std::string_view s) const {
		assert(s.size() <= std::size_t(std::numeric_limits<int>::max()));
		Match best;
		int v = root, matched = 0;
		for (int i = 0; i < static_cast<int>(s.size()); ++i) {
			int next = transition(v, s[i]);
			while (v && next == -1) {
				v = nodes_[v].link;
				matched = nodes_[v].length;
				next = transition(v, s[i]);
			}
			if (next == -1) { v = root; matched = 0; continue; }
			v = next, ++matched;
			if (matched > best.length) best = {nodes_[v].first_end - matched, i + 1 - matched, matched};
		}
		return best;
	}
};

#endif
Back to top page