cp-library

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

View the Project on GitHub qandori/cp-library

:heavy_check_mark: 線形漸化式の復元(Berlekamp–Massey)

fps/berlekamp_massey.hpp · GitHub で見る

共通ページ

有限数列と矛盾しない定数係数線形漸化式のうち、位数が最小のものを 1 つ求めます。

#include "lib/fps/berlekamp_massey.hpp"
#include "lib/math/mint.hpp"

std::vector<mint> a{0, 1, 1, 2, 3, 5, 8, 13};
auto c = berlekamp_massey(a); // {1,1}
操作 内容・返り値 計算量
berlekamp_massey(a) 最小位数の漸化式の係数を vector<T> で返す \(O(N^2)\)

返り値の長さを \(d\) とすると、入力の全ての \(d\le i<N\) について次が成り立ちます。

\[\displaystyle a_i=\sum_{j=0}^{d-1}c_j a_{i-j-1}\]

c[0] は 1 つ前の項の係数です。法 998244353 なら線形漸化式の第 k 項にそのまま渡せます。次は上の a・c を使う例です。初期値は a の先頭 c.size() 項を渡してください。

#include "lib/fps/linear_recurrence.hpp"

std::vector<mint> initial(a.begin(), a.begin() + c.size());
auto value = linear_recurrence(initial, c, 10); // 55

T は正確な四則演算と 0 の等値判定ができる体の型を使います。mint、atcoder::modint1000000007 などの素数 modint、有理数型に対応します。整数型や合成数 mod の modint は対象外です。入力長は INT_MAX 未満としてください。

空の入力や全て 0 の数列には空の係数列を返します。位数に必要な末尾の 0 は残します。例えば {1,0,0} に対する返り値は {0} です。

解が複数ある場合はそのうち 1 つを返します。元の無限数列の位数が \(D\) 以下と分かっていれば、先頭 \(2D\) 項以上から後続項も定まります。

入力長を \(N\)、返す位数を \(d\) とし、体の演算を \(O(1)\) とすると、時間は \(O(Nd)\)、最悪 \(O(N^2)\)、追加空間は \(O(N)\) です。全て 0 の数列は \(O(N)\) 時間です。

参考:AtCoder Algorithm Lectures:線形漸化式の復元

依存ライブラリ

検証コード

コード

#ifndef BERLEKAMP_MASSEY_HPP
#define BERLEKAMP_MASSEY_HPP

#include "../components/field_arithmetic.hpp"
#include <cassert>
#include <limits>
#include <type_traits>
#include <vector>

namespace berlekamp_massey_detail {
template <class Ops>
std::vector<typename Ops::V> restore(const std::vector<typename Ops::V> &sequence) {
	using V = typename Ops::V;
	const int n = int(sequence.size());
	std::vector<V> reversed(sequence.rbegin(), sequence.rend()), current{V(1)}, previous{V(1)}, saved;
	current.reserve(sequence.size() + 1);
	previous.reserve(sequence.size() + 1);
	saved.reserve(sequence.size() + 1);
	int order = 0, shift = 1;
	V inverse = 1;
	for (int i = 0; i < n; ++i) {
		V discrepancy = Ops::dot(current.data(), reversed.data() + n - 1 - i, int(current.size()));
		if (Ops::zero(discrepancy, 0)) {
			++shift;
			continue;
		}
		bool grow = order <= i / 2;
		if (grow) saved.assign(current.begin(), current.end());
		V factor = Ops::neg(Ops::mul(discrepancy, inverse));
		if (current.size() < previous.size() + shift) current.resize(previous.size() + shift);
		Ops::axpy(current.data() + shift, previous.data(), int(previous.size()), factor);
		if (grow) {
			order = i + 1 - order;
			previous.swap(saved);
			inverse = Ops::inv(discrepancy);
			shift = 1;
		} else ++shift;
	}
	current.resize(std::size_t(order) + 1);
	current.erase(current.begin());
	for (auto &x : current) x = Ops::neg(x);
	return current;
}
}

template <class T>
std::vector<T> berlekamp_massey(const std::vector<T> &sequence) {
	static_assert(!std::is_integral_v<T> && !std::is_same_v<T, __int128_t> && !std::is_same_v<T, __uint128_t>);
	assert(sequence.size() < std::size_t(std::numeric_limits<int>::max()));
	if constexpr (field_detail::small_modint<T>) {
		std::vector<std::uint32_t> values(sequence.size());
		for (std::size_t i = 0; i < sequence.size(); ++i) values[i] = sequence[i].val();
		auto coefficients = berlekamp_massey_detail::restore<field_detail::ModularField<T>>(values);
		std::vector<T> result;
		result.reserve(coefficients.size());
		for (auto x : coefficients) result.push_back(T::raw(x));
		return result;
	} else return berlekamp_massey_detail::restore<field_detail::Field<T>>(sequence);
}

#endif
Back to top page