atcoder/lib/math/
mod_int.rs

1//! コンパイル時に法を指定する剰余整数です。
2
3use std::fmt::{self, Display};
4use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};
5
6/// コンパイル時定数 `MOD` を法とする剰余整数です。
7///
8/// 構築時に値を剰余で正規化し、保持値は常に `0 <= value < MOD` です。`MOD` は正である必要が
9/// あります。`pow` は非負整数の指数に使います。`inv` と `/` は `MOD` が素数で、対象の値が
10/// 法 `MOD` で `0` でない場合にのみ数学的な乗法逆元として機能します。
11/// 加算・減算・乗算・累乗の途中計算は `i64` で行うため、中間値が `i64` に収まる `MOD` を
12/// 選んでください。乗算については、少なくとも `(MOD - 1) * (MOD - 1)` が `i64::MAX` 以下に
13/// なる法であることを確認してください。
14///
15/// # Panics
16///
17/// `MOD <= 0` の場合は構築時に panic します。
18///
19/// # Examples
20///
21/// ```
22/// use atcoder::math::mod_int::Mod998;
23///
24/// let value = Mod998::new(-1);
25/// assert_eq!(value.value(), 998_244_352);
26/// assert_eq!((value + Mod998::new(2)).value(), 1);
27/// ```
28#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
29pub struct ModInt<const MOD: i64> {
30    value: i64,
31}
32
33impl<const MOD: i64> ModInt<MOD> {
34    /// `value` を法 `MOD` で正規化した値を作ります。負の値も `0..MOD` に収めます。
35    /// 計算量は `O(1)` です。
36    ///
37    /// # Panics
38    ///
39    /// `MOD <= 0` の場合に panic します。
40    pub fn new(value: i64) -> Self {
41        assert!(MOD > 0, "modulus must be positive");
42        Self {
43            value: value.rem_euclid(MOD),
44        }
45    }
46
47    /// 正規化済みの値（`0..MOD`）を返します。
48    pub const fn value(self) -> i64 {
49        self.value
50    }
51
52    /// 正規化済みの値（`0..MOD`）を参照経由で返します。
53    ///
54    /// [`value`](Self::value) と同じです。
55    pub const fn val(&self) -> i64 {
56        self.value
57    }
58
59    /// フェルマーの小定理に基づき、`self` の乗法逆元候補を返します。
60    ///
61    /// `MOD` が素数かつ `self.value() != 0` の場合に、数学的な逆元になることを保証します。
62    /// 素数判定やゼロの検査は行いません。`MOD == 1` では指数が負になるため panic します。
63    /// それ以外の条件外の入力でも逆元としての正しさは保証しないため、呼び出し側で条件を保証してください。
64    /// 計算量は `O(log MOD)` です。
65    pub fn inv(self) -> Self {
66        self.pow(MOD - 2)
67    }
68
69    /// `self` の `exponent` 乗を法 `MOD` で返します。
70    ///
71    /// `exponent` は非負である必要があります。計算量は `O(log exponent)` です。
72    /// 乗算中間値を `i64` で計算するため、想定法に対してその値が範囲内に収まることが必要です。
73    ///
74    /// # Panics
75    ///
76    /// `exponent < 0` の場合に panic します。
77    pub fn pow(self, mut exponent: i64) -> Self {
78        assert!(exponent >= 0, "exponent must be non-negative");
79        let mut result = 1_i64;
80        let mut base = self.value;
81        while exponent > 0 {
82            if exponent & 1 == 1 {
83                result = result * base % MOD;
84            }
85            base = base * base % MOD;
86            exponent >>= 1;
87        }
88        Self::new(result)
89    }
90}
91
92impl<const MOD: i64> Display for ModInt<MOD> {
93    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
94        self.value.fmt(formatter)
95    }
96}
97
98impl<const MOD: i64> Add for ModInt<MOD> {
99    type Output = Self;
100    fn add(self, rhs: Self) -> Self::Output {
101        Self::new(self.value + rhs.value)
102    }
103}
104
105impl<const MOD: i64> Sub for ModInt<MOD> {
106    type Output = Self;
107    fn sub(self, rhs: Self) -> Self::Output {
108        Self::new(self.value - rhs.value)
109    }
110}
111
112impl<const MOD: i64> Mul for ModInt<MOD> {
113    type Output = Self;
114    fn mul(self, rhs: Self) -> Self::Output {
115        Self::new(self.value * rhs.value)
116    }
117}
118
119impl<const MOD: i64> Div for ModInt<MOD> {
120    type Output = Self;
121    #[allow(clippy::suspicious_arithmetic_impl)]
122    fn div(self, rhs: Self) -> Self::Output {
123        self * rhs.inv()
124    }
125}
126
127impl<const MOD: i64> AddAssign for ModInt<MOD> {
128    fn add_assign(&mut self, rhs: Self) {
129        *self = *self + rhs;
130    }
131}
132
133impl<const MOD: i64> SubAssign for ModInt<MOD> {
134    fn sub_assign(&mut self, rhs: Self) {
135        *self = *self - rhs;
136    }
137}
138
139impl<const MOD: i64> MulAssign for ModInt<MOD> {
140    fn mul_assign(&mut self, rhs: Self) {
141        *self = *self * rhs;
142    }
143}
144
145impl<const MOD: i64> DivAssign for ModInt<MOD> {
146    fn div_assign(&mut self, rhs: Self) {
147        *self = *self / rhs;
148    }
149}
150
151/// 法 `998_244_353` の剰余整数です（競技プログラミングでよく使われる素数法）。
152pub type Mod998 = ModInt<998_244_353>;
153/// 法 `1_000_000_007` の剰余整数です（競技プログラミングでよく使われる素数法）。
154pub type Mod107 = ModInt<1_000_000_007>;
155
156#[cfg(test)]
157mod tests {
158    use super::Mod998;
159
160    #[test]
161    fn normalizes_and_calculates() {
162        assert_eq!(Mod998::new(-1).value(), 998_244_352);
163        assert_eq!((Mod998::new(2) / Mod998::new(2)).value(), 1);
164        assert_eq!(Mod998::new(2).pow(10).value(), 1024);
165    }
166}
atcoder/lib/math/mod_int.rs

atcoder/lib/math/
mod_int.rs
