Online Convolution（高速版）

convolution/online_convolution.hpp · GitHub で見る

全体サイト（256 KiB向け） · 高速実装

軽量実装（64 KiB）はこちら

#include "lib/convolution/online_convolution.hpp"

OnlineConvolution<> product(3);
mint c0 = product.push(2, 4);
mint c1 = product.push(3, 5);
mint known = product.pending();
mint c2 = product.push(7, 11);

a・b の係数を定数項から同時に 1 項ずつ追加し、積の次の係数を返します。上の例では c0=8、c1=22、known=15、c2=65 です。将来の入力は不要で、確定した結果から次の入力を決められます。

関数・型	内容・返り値	計算量
OnlineConvolution<T = mint>	法 998244353 のオンライン畳み込み。係数型 T は mod()・val()・raw() を持つ modint	—（型）
OnlineConvolution(max_size = 1 << 23)	追加できる項数の上限を指定して構築。最初から上限分を確保することはない	\(O(1)\)
push(x,y)	次の係数対 a[n]=x・b[n]=y を追加し、積の対応する係数を T で返す。追加前の項数を n とすると \(\displaystyle \sum_{i=0}^n a_i b_{n-i}\)	償却 \(O(\log^2 n)\)、最悪 \(O(n\log n)\)
pending()	次の係数のうち、追加済みの入力だけで決まる部分を T で返す	\(O(1)\)
operator[](i)	確定済みの積の第 i 係数を T で返す。0 <= i < size()	\(O(1)\)
size()	追加済みの係数対の数	\(O(1)\)
max_size()	構築時に指定した項数の上限	\(O(1)\)
empty()	係数対をまだ追加していなければ true	\(O(1)\)
reserve(n)	入力列と結果の領域を事前確保。0 <= n <= max_size()。上限は変わらない	\(O(n)\)
clear()	入力・結果・変換済みの先頭区間を消去して再利用。上限は変わらない	\(O(n)\)

0 <= max_size <= 2^23 とし、push()・pending() は size() < max_size() のときに呼びます。必要な項数が分かる場合は、構築時に指定すると最後の不要な先読み計算を省けます。長さの異なる有限列を最後まで畳み込む場合は、上限を a.size()+b.size()-1 とし、それぞれの列の末尾に 0 を追加します。

追加済みの項数を \(n\) とすると、pending() は \(\displaystyle \sum_{i=1}^{n-1} a_i b_{n-i}\) です。\(n=0,1\) では 0 を返します。次の push(x,y) の返り値は、\(n>0\) なら \(\mathrm{pending}()+xb_0+a_0y\)、\(n=0\) なら \(xy\) になります。pending() は状態を変更しません。

係数の法は 998244353 に固定しています。mint を既定とし、同じ法の別の modint でも使えます。

項数を \(n\) として、push() を合計 \(n\) 回行う時間は \(O(n\log^2 n)\)、1 回の最悪時間は \(O(n\log n)\)、使用空間は \(O(n)\) です。構築・size()・max_size()・empty()・operator[]・pending() は \(O(1)\) です。reserve(r) は \(O(n)\) 時間・\(O(r)\) 空間、clear() は \(O(n)\) 時間以内です。事前確保した領域と共有 NTT キャッシュは別途保持されます。

依存ライブラリ

• modint998244353 の共通剰余演算
• NTT（998244353）（高速版）
• modint998244353

検証コード

• test/online_convolution.test.cpp
• test/online_convolution_exp.test.cpp
• test/online_convolution_inv.test.cpp

コード

#ifndef ONLINE_CONVOLUTION_HPP
#define ONLINE_CONVOLUTION_HPP

#include "../components/ntt.hpp"
#include <algorithm>
#include <bit>
#include <cassert>
#include <type_traits>
#include <vector>

template <class T = mint>
class OnlineConvolution {
	static_assert(T::mod() == 998244353);
	using M = mint;
	using NTT = fps_detail::NTT;
	using V = std::vector<M>;
	struct Prefix { V a, b; };
	int limit;
	V a, b, c, f, g;
	std::vector<Prefix> prefixes;
	static M import(T x) {
		if constexpr (std::is_same_v<T, M>) return x;
		else return M::raw(x.val());
	}
	static T export_value(M x) {
		if constexpr (std::is_same_v<T, M>) return x;
		else return T::raw(x.val());
	}
	static NTT::U *words(V &v) { return reinterpret_cast<NTT::U *>(v.data()); }
	static const NTT::U *words(const V &v) { return reinterpret_cast<const NTT::U *>(v.data()); }
	static void forward(V &v) {
		NTT &plan = fps_detail::ntt_plan(std::max(1, int(v.size()) / 2));
		plan.fft<false>(words(v), v.size());
	}
	static void inverse(V &v) {
		int n = v.size();
		NTT &plan = fps_detail::ntt_plan(std::max(1, n / 2));
		plan.fft<true>(words(v), n);
		plan.scale(words(v), n, M::mod() - (M::mod() - 1) / n);
	}
	void extend(int q, int d) {
		int end = std::min(q + d, limit);
		if (int(c.size()) < end) c.resize(end);
		if (d < 8) {
			for (int k = q; k < end; ++k) {
				M sum = 0;
				if (q == d) {
					for (int i = k - q + 1; i < q; ++i) sum += a[i] * b[k - i];
				} else {
					for (int i = q - d; i < q; ++i) sum += a[i] * b[k - i] + b[i] * a[k - i];
				}
				c[k] += sum;
			}
			return;
		}
		if (q == d) {
			int level = std::countr_zero(unsigned(q));
			prefixes.resize(level + 1);
			auto &p = prefixes[level];
			p.a.assign(a.begin(), a.end());
			p.b.assign(b.begin(), b.end());
			forward(p.a); forward(p.b);
			f.resize(q);
			NTT::product(words(f), words(p.a), words(p.b), q);
			inverse(f);
			for (int k = q; k < end; ++k) c[k] += f[k - q] - c[k - q];
		} else {
			int n = 2 * d, level = std::countr_zero(unsigned(n));
			const auto &p = prefixes[level];
			f.assign(n, M(0)); g.assign(n, M(0));
			std::copy_n(a.begin() + q - d, d, f.begin());
			std::copy_n(b.begin() + q - d, d, g.begin());
			forward(f); forward(g);
			NTT::product(words(f), words(f), words(p.b), n);
			NTT::product(words(g), words(g), words(p.a), n);
			for (int i = 0; i < n; ++i) f[i] += g[i];
			inverse(f);
			for (int k = q; k < end; ++k) c[k] += f[d + k - q];
		}
	}
public:
	explicit OnlineConvolution(int max_size = 1 << 23) : limit(max_size) { assert(0 <= limit && limit <= (1 << 23)); }
	int size() const { return a.size(); }
	int max_size() const { return limit; }
	bool empty() const { return a.empty(); }
	void reserve(int n) {
		assert(0 <= n && n <= limit);
		a.reserve(n); b.reserve(n); c.reserve(n);
	}
	void clear() { a.clear(); b.clear(); c.clear(); f.clear(); g.clear(); prefixes.clear(); }
	T operator[](int i) const { assert(0 <= i && i < size()); return export_value(c[i]); }
	T pending() const { assert(size() < limit); return export_value(size() < int(c.size()) ? c[size()] : M(0)); }
	T push(T x, T y) {
		int n = size();
		assert(n < limit);
		a.push_back(import(x)); b.push_back(import(y));
		if (n >= int(c.size())) c.resize(n + 1);
		c[n] += a[n] * b[0];
		if (n) c[n] += b[n] * a[0];
		int q = n + 1;
		if (q < limit) extend(q, q & -q);
		return export_value(c[n]);
	}
};

#endif

Back to top page