二項係数 prefix sum のオンライン処理部品

components/binomial_prefix_online.hpp · GitHub で見る

共通ページ

二項係数の prefix sum の通常版・軽量版で共有する内部実装です。利用するときは公開ヘッダーを include してください。

binomial_prefix_detail::Online<Mint,Table> は、素数 mod の要素型 Mint と階乗・逆階乗・2 の累乗を生成する Table を使います。二次元格子に正規化した prefix を保存し、近い格子点から積和で補正します。max_n は非負かつ INT_MAX と法より小さく、block_size は非負にします。

操作	内容・返り値	計算量
Online(max_n=0,block_size=0)	上限とブロック幅を指定して構築。幅 0 は平方根を基準に決定	\(O(N+N^2/B)\)
build(max_n,block_size=0)	前計算を作り直す。返り値は void	\(O(N+N^2/B)\)
query(n,k)	右端を含まない prefix を Mint で返す	\(O(B)\)
prod(n,l,r)	l <= r の区間和を Mint で返す	\(O(B)\)
max_n()	上限を int で返す	\(O(1)\)
block_size()	ブロック幅を int で返す	\(O(1)\)
table_size()	格子に保存した Mint の個数を std::size_t で返す	\(O(1)\)

n は前計算した範囲内で指定し、k・l・r は std::int64_t です。負の添字と n より大きい添字の二項係数は 0 として扱います。上限 0 にも対応します。

上限を \(N\)、ブロック幅を \(B\) とすると、構築は \(O(N+N^2/B)\) 時間、保存量は \(O(N+N^2/B^2)\)、取得は \(O(B)\) 時間です。

依存ライブラリ

• mint の積和部品
• modint998244353

検証コード

• test/binomial_prefix_sum.test.cpp
• test/binomial_prefix_sum.test.cpp
• test/binomial_prefix_sum_light.test.cpp
• test/binomial_prefix_sum_light.test.cpp
• test/binomial_prefix_sum_online.test.cpp
• test/binomial_prefix_sum_online.test.cpp
• test/binomial_prefix_sum_online_light.test.cpp
• test/binomial_prefix_sum_online_light.test.cpp

コード

#ifndef BINOMIAL_PREFIX_ONLINE_HPP
#define BINOMIAL_PREFIX_ONLINE_HPP

#include "mint_inner_product.hpp"
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <vector>

namespace binomial_prefix_detail {
template<class Mint, class Table> class Online {
	Table table_;
	int max_n_ = 0, block_ = 1;
	std::vector<Mint> scaled_factorial_, data_;
	void check([[maybe_unused]] int n) const { assert(0 <= n && n <= max_n_); }
	static std::size_t offset(int row) { return std::size_t(row) * (std::size_t(row) + 1) / 2; }
	template<bool Reverse> static Mint dot(const Mint* a, const Mint* b, int size) {
		if constexpr (std::is_same_v<Mint, mint>) {
			return mint_detail::inner_product<Reverse>(a, b, size);
		} else {
			Mint sum(0);
			for (int i = 0; i < size; ++i) sum += a[i] * b[Reverse ? -i : i];
			return sum;
		}
	}
	Mint range(int n, int l, int r) const {
		return table_.factorial[n] * dot<true>(table_.inverse_factorial.data() + l, table_.inverse_factorial.data() + n - l, r - l);
	}

  public:
	explicit Online(int max_n = 0, int block_size = 0) { build(max_n, block_size); }
	void build(int max_n, int block_size = 0) {
		assert(max_n >= 0 && block_size >= 0);
		table_.build(max_n);
		max_n_ = max_n;
		block_ = block_size ? block_size : std::max(1, int(std::sqrt(double(max_n) + 1)));
		if (max_n < 2) { scaled_factorial_.clear(); data_.clear(); return; }
		scaled_factorial_.resize(std::size_t(max_n) + 1);
		Mint inverse_power = table_.power2[max_n].inv() * table_.half;
		for (int i = max_n; i >= 0; --i) {
			scaled_factorial_[i] = table_.factorial[i] * inverse_power;
			inverse_power += inverse_power;
		}
		int rows = max_n / block_;
		data_.resize(offset(rows) + std::size_t(rows) + 1);
		for (int row = 0; row <= rows; ++row) {
			int n = row * block_;
			Mint factor = scaled_factorial_[n] + scaled_factorial_[n], sum(0);
			auto base = offset(row);
			for (int column = 0; column <= row / 2; ++column) {
				int k = column * block_;
				if (column) sum += dot<true>(table_.inverse_factorial.data() + k - block_, table_.inverse_factorial.data() + n - k + block_, block_);
				Mint value = sum * factor;
				data_[base + row - column] = Mint(1) - value;
				data_[base + column] = value + table_.inverse_factorial[k] * table_.inverse_factorial[n - k] * factor;
			}
		}
	}
	int max_n() const { return max_n_; }
	int block_size() const { return block_; }
	std::size_t table_size() const { return data_.size(); }
	Mint query(int n, std::int64_t k) const {
		check(n);
		if (k <= 0) return Mint(0);
		if (k >= std::int64_t(n) + 1) return table_.power2[n];
		bool reflect = k > (n + 1) / 2;
		int m = reflect ? n + 1 - int(k) : int(k);
		Mint value(1);
		if (m != 1) {
			if (m == (n + 1) / 2) value = table_.center(n);
			else {
				int row = std::min(int((std::int64_t(n) + block_ / 2) / block_), max_n_ / block_);
				int column = std::min(int((std::int64_t(m - 1) + block_ / 2) / block_), row);
				int start = row * block_, anchor = column * block_;
				int center = (n + 1) / 2;
				auto distance = std::int64_t(std::abs(n - start)) + std::abs(m - 1 - anchor);
				if (std::min(m - 1, center - m) <= distance) {
					value = m - 1 <= center - m ? Mint(1) + range(n, 1, m) : table_.center(n) - range(n, m, center);
				} else {
					value = data_[offset(row) + column];
					if (start < n) value -= table_.inverse_factorial[anchor] * dot<false>(scaled_factorial_.data() + start, table_.inverse_factorial.data() + start - anchor, n - start);
					else if (n < start) {
						int first = std::max(n, anchor);
						value += table_.inverse_factorial[anchor] * dot<false>(scaled_factorial_.data() + first, table_.inverse_factorial.data() + first - anchor, start - first);
					}
					value *= table_.power2[n];
					if (anchor < m - 1) value += range(n, anchor + 1, m);
					else if (m - 1 < anchor) value -= range(n, m, std::min(anchor, n) + 1);
				}
			}
		}
		return reflect ? table_.power2[n] - value : value;
	}
	Mint prod(int n, std::int64_t l, std::int64_t r) const { assert(l <= r); return query(n, r) - query(n, l); }
};
}

#endif

Back to top page