競技プログラミング用 C++ ライブラリ
math/modular_arithmetic.hpp · GitHub で見る
共通ページ
#include "lib/math/modular_arithmetic.hpp"
#include <cassert>
int main() {
assert(pow_mod(-2, 10, 1000) == 24);
assert(inv_mod(3, 10) == 7);
}
| 関数 | 引数・返り値 |
|---|---|
pow_mod(x, n, m) |
x, n は std::int64_t、m は std::uint32_t。n >= 0, m >= 1 で \(x^n\bmod m\) を std::uint32_t で返す |
inv_mod(x, m) |
両引数は std::int64_t、m >= 1。\(xy\equiv1\pmod m\) を満たす 0 <= y < m を std::int64_t で返す |
どちらも素数・合成数の法と負の x に対応します。法 1 の返り値は 0 です。0 の 0 乗は 1 を法で正規化した値です。不正な法・負の指数は std::invalid_argument、逆元が存在しない場合は std::domain_error です。inv_mod は定数式でも利用できます。
pow_mod は \(O(\log n)\) 時間、inv_mod は \(O(\log m)\) 時間、追加空間はいずれも \(O(1)\) です。指数 0 は定数時間です。
#ifndef MODULAR_ARITHMETIC_HPP
#define MODULAR_ARITHMETIC_HPP
#include "linear_congruence.hpp"
#include "../components/barrett32.hpp"
#include <stdexcept>
inline std::uint32_t pow_mod(std::int64_t x, std::int64_t n, std::uint32_t m) {
if (!m || n < 0) throw std::invalid_argument("pow_mod arguments");
std::int64_t r = x % std::int64_t(m);
if (r < 0) r += m;
return Barrett32(m).pow(std::uint32_t(r), std::uint64_t(n));
}
constexpr std::int64_t inv_mod(std::int64_t x, std::int64_t m) {
auto result = linear_congruence(x, 1, m);
if (!result) throw std::domain_error("inverse does not exist");
return result->residue;
}
#endif
#line 1 "math/modular_arithmetic.hpp"
#line 1 "math/linear_congruence.hpp"
#include <cstdint>
#include <optional>
#include <stdexcept>
#include <utility>
struct LinearCongruenceSolution {
std::int64_t residue, period, count;
bool operator==(const LinearCongruenceSolution &) const = default;
};
constexpr std::optional<LinearCongruenceSolution> linear_congruence(std::int64_t a, std::int64_t b, std::int64_t m) {
if (m <= 0) throw std::invalid_argument("linear_congruence modulus");
a %= m;
b %= m;
if (a < 0) a += m;
if (b < 0) b += m;
std::int64_t r = m, s = a, u = 0, v = 1;
while (s) {
std::int64_t q = r / s;
r -= q * s;
u -= q * v;
std::swap(r, s);
std::swap(u, v);
}
if (b % r) return std::nullopt;
std::int64_t period = m / r;
if (u < 0) u += period;
std::int64_t residue = __int128_t(u) * (b / r) % period;
return LinearCongruenceSolution{residue, period, r};
}
#line 1 "components/barrett32.hpp"
#include <cassert>
#line 6 "components/barrett32.hpp"
struct Barrett32 {
std::uint32_t mod;
std::uint64_t reciprocal;
explicit Barrett32(std::uint32_t m) : mod(m), reciprocal(m > 1 ? std::uint64_t((__uint128_t(1) << 64) / m) : 0) {
assert(m > 0);
}
std::uint32_t reduce(std::uint64_t x) const {
if (mod == 1) return 0;
std::uint64_t q = (__uint128_t(x) * reciprocal) >> 64, r = x - q * mod;
return r < mod ? r : r - mod;
}
std::uint32_t mul(std::uint32_t a, std::uint32_t b) const { return reduce(std::uint64_t(a) * b); }
std::uint32_t pow(std::uint32_t a, std::uint64_t e) const {
std::uint32_t r = mod == 1 ? 0 : 1;
for (; e; e >>= 1, a = mul(a, a)) if (e & 1) r = mul(r, a);
return r;
}
};
#line 7 "math/modular_arithmetic.hpp"
inline std::uint32_t pow_mod(std::int64_t x, std::int64_t n, std::uint32_t m) {
if (!m || n < 0) throw std::invalid_argument("pow_mod arguments");
std::int64_t r = x % std::int64_t(m);
if (r < 0) r += m;
return Barrett32(m).pow(std::uint32_t(r), std::uint64_t(n));
}
constexpr std::int64_t inv_mod(std::int64_t x, std::int64_t m) {
auto result = linear_congruence(x, 1, m);
if (!result) throw std::domain_error("inverse does not exist");
return result->residue;
}