定数時間 mod pow

math/mod_pow.hpp · GitHub で見る

共通ページ · 軽量実装なし

#include "lib/math/mod_pow.hpp"

ModPow power(17);
auto a = power.pow(3, 5); // 5
auto b = power.inv(7); // 5
auto k = power.log(5); // 5
auto c = power.pow_root(k); // 5

操作	内容・返り値	計算量
ModPow(prime = 998244353,root = 0)	素数を固定し、累乗・逆元・離散対数用の表を前計算。root = 0 なら原始根を自動探索	原始根指定時: 期待 \(O(p)\)、自動探索は本文参照
mod()	法を uint32_t で返す	\(O(1)\)
primitive_root()	使用している原始根を uint32_t で返す	\(O(1)\)
pow(a,exponent)	\(a^{\mathrm{exponent}}\bmod p\)。\(0\le a<p\)、指数は非負の uint64_t。\(0^0=1\)	\(O(1)\)
inv(x)	\(x^{-1}\bmod p\)。\(1\le x<p\)	\(O(1)\)
log(x)	使用している原始根 \(g\) に対する \(g^k=x\pmod p\) の最小の非負整数 \(k\)。\(1\le x<p\)	\(O(1)\)
pow_root(exponent)	\(g^{\mathrm{exponent}}\bmod p\)。指数は非負の uint64_t	\(O(1)\)

法 \(p\) は \(2\le p<2^{31}\) の素数です。底と指数はクエリごとに変更できます。法と原始根は実行時に指定でき、root を省略すると最小の原始根を探索します。返り値はすべて uint32_t です。底は正規化済みの値を渡します。

同じ法で多数の累乗・逆元を計算するときに、オブジェクトを一度作って使い回します。例の法 17 では原始根は 3、log(5) は「3 を何乗すると 5 になるか」を返します。通常の pow の底は原始根以外でも指定できます。大きい法では前計算に時間とメモリを使うため、少数回だけの累乗には mint::pow なども使えます。

pow は原始根に対する離散対数に変換し、指数を \(p-1\) で縮約してから表を参照します。有理数近似で参照範囲を小さくするため、法の大きさに比例する全対数表は持ちません。逆元だけが必要なら ModInv を使うと、離散対数の前計算を省けます。mint と組み合わせる場合は mint::raw(power.pow(a.val(),k)) などとします。

離散対数の前計算では小さい素数の対数を利用した乱択を最大 64 回試し、失敗した場合は BSGS に切り替えます。乱数 seed は法と原始根から固定され、結果と処理は再現可能です。クエリには乱択もループも使いません。

法や底を毎回変更する場合や、合成数の法を扱う場合は discrete_log を使えます。

前計算の空間は \(O(p^{2/3})\)、原始根が与えられた場合の時間はハッシュ探索を期待 \(O(1)\) として期待 \(O(p)\) です。原始根の自動探索には、\(p-1\) の試し割りと候補ごとの累乗判定が追加されます。すべてのクエリは \(O(1)\) 時間です。

依存ライブラリ

• 素数判定の基数表
• 素数判定の共通演算
• 定数時間 mod 演算の内部部品
• Montgomery 64
• Montgomery 32
• SplitMix64
• 素数判定（高速版）
• 定数時間 mod inv

検証コード

• test/mod_pow_log.test.cpp

コード

#ifndef MOD_POW_HPP
#define MOD_POW_HPP

#include "mod_inv.hpp"
#include "../components/splitmix64.hpp"
#include <initializer_list>

class ModPow : public ModInv {
	mod_fast_detail::Reducer order_;
	std::uint32_t root_ = 1, shift_, mask_;
	std::vector<std::uint32_t> low_, high_, logarithm_;
	std::uint32_t power_(std::uint32_t e) const { return modulus_.mul(low_[e & mask_], high_[e >> shift_]); }
	void build_log_() {
		std::uint32_t p = mod(), limit = fractions_.limit(), order = p - 1;
		logarithm_.resize(inverse_.size());
		if (!center_) {
			std::uint32_t x = 1;
			for (std::uint32_t i = 0; i < order; ++i, x = modulus_.mul(x, root_)) logarithm_[x] = i;
			return;
		}
		std::vector<std::uint16_t> factor(limit + 1);
		for (std::uint32_t q = 2; std::uint64_t(q) * q <= limit; ++q) if (!factor[q]) {
			for (std::uint32_t i = q * q; i <= limit; i += q) if (!factor[i]) factor[i] = q;
		}
		std::uint32_t step = std::min(order, 4 * (mask_ + 1)), capacity = std::bit_ceil(2 * step);
		std::vector<std::uint64_t> baby(capacity);
		auto hash = [=](std::uint32_t x) {
			x ^= x >> 16, x *= 0x7feb352dU, x ^= x >> 15;
			return x & (capacity - 1);
		};
		std::uint32_t value = 1;
		for (std::uint32_t i = 0; i < step; ++i, value = modulus_.mul(value, root_)) {
			auto k = hash(value);
			while (baby[k]) k = (k + 1) & (capacity - 1);
			baby[k] = (std::uint64_t(value) << 32) | i;
		}
		std::uint32_t giant = power_(order - step);
		auto bsgs = [&](std::uint32_t x) {
			for (std::uint32_t i = 0; i < order; i += step, x = modulus_.mul(x, giant)) {
				auto k = hash(x);
				while (baby[k] && std::uint32_t(baby[k] >> 32) != x) k = (k + 1) & (capacity - 1);
				if (baby[k]) {
					std::uint32_t result = i + std::uint32_t(baby[k]);
					return result < order ? result : result - order;
				}
			}
			assert(false);
			return std::uint32_t(0);
		};
		SplitMix64 rng((std::uint64_t(p) << 32) | root_);
		for (std::uint32_t i = 2; i <= limit; ++i) {
			std::uint32_t result = 0;
			if (factor[i]) result = order_.reduce(std::uint64_t(logarithm_[center_ + factor[i]]) + logarithm_[center_ + i / factor[i]]);
			else if (std::uint64_t(i) * i > p) {
				result = order_.reduce(std::uint64_t(logarithm_[center_ + p % i]) + order / 2 + order - logarithm_[center_ + p / i]);
			} else {
				bool found = false;
				if (i >= 100) for (int attempt = 0; attempt < 64 && !found; ++attempt) {
					std::uint32_t k = order_.reduce(rng()), x = modulus_.mul(i, power_(k));
					auto [a, b] = fractions_.fraction(x);
					std::uint64_t sum = order - k + (a < 0 ? order / 2 : 0);
					auto accumulate = [&](std::uint32_t value, bool subtract) {
						while (value > 1) {
							std::uint32_t q = factor[value] ? factor[value] : value;
							if (q >= i) return false;
							sum += subtract ? order - logarithm_[center_ + q] : logarithm_[center_ + q];
							value /= q;
						}
						return true;
					};
					if (accumulate(a < 0 ? -a : a, false) && accumulate(b, true)) found = true, result = order_.reduce(sum);
				}
				if (!found) result = bsgs(i);
			}
			logarithm_[center_ + i] = result;
		}
		for (std::uint32_t i = 1; i <= limit; ++i)
			logarithm_[center_ - i] = order_.reduce(std::uint64_t(logarithm_[center_ + i]) + order / 2);
	}

  public:
	explicit ModPow(std::uint32_t prime = 998244353, std::uint32_t root = 0) : ModInv(prime), order_(prime - 1) {
		std::uint32_t remaining = prime - 1;
		std::vector<std::uint32_t> factors;
		for (std::uint32_t q = 2; std::uint64_t(q) * q <= remaining; ++q) if (remaining % q == 0) {
			factors.push_back(q);
			do { remaining /= q; } while (remaining % q == 0);
		}
		if (remaining > 1) factors.push_back(remaining);
		auto valid = [&](std::uint32_t g) {
			return 0 < g && g < prime && std::all_of(factors.begin(), factors.end(), [&](auto q) { return modulus_.pow(g, (prime - 1) / q) != 1; });
		};
		if (root) { assert(valid(root)); root_ = root; }
		else while (!valid(root_)) ++root_;
		shift_ = (std::bit_width(prime - 2) + 1) / 2;
		std::uint32_t width = std::uint32_t(1) << shift_;
		mask_ = width - 1;
		low_.resize(width), high_.resize((prime - 2) / width + 1);
		low_[0] = high_[0] = 1;
		for (std::uint32_t i = 1; i < width; ++i) low_[i] = modulus_.mul(low_[i - 1], root_);
		std::uint32_t stride = modulus_.mul(low_.back(), root_);
		for (std::uint32_t i = 1; i < high_.size(); ++i) high_[i] = modulus_.mul(high_[i - 1], stride);
		build_log_();
	}
	std::uint32_t primitive_root() const { return root_; }
	std::uint32_t log(std::uint32_t x) const {
		assert(0 < x && x < mod());
		if (x <= fractions_.limit()) return logarithm_[center_ + x];
		auto [a, b] = fractions_.fraction(x);
		std::uint32_t result = logarithm_[center_ + a] + (mod() - 1) - logarithm_[center_ + b];
		return result < mod() - 1 ? result : result - (mod() - 1);
	}
	std::uint32_t pow_root(std::uint64_t exponent) const { return power_(order_.reduce(exponent)); }
	std::uint32_t pow(std::uint32_t a, std::uint64_t exponent) const {
		assert(a < mod());
		if (!a) return exponent == 0 ? 1 : 0;
		return power_(order_.mul(log(a), order_.reduce(exponent)));
	}
};

#endif

Back to top page