atcoder/lib/math/
fast_math.rs

1//! 組み込み整数型向けの基本演算です。
2
3/// 高速累乗、剰余逆元、最大公約数、最小公倍数を追加する拡張トレイトです。
4///
5/// `i32`、`i64`、`isize`、`u32`、`u64`、`usize` に実装されています。演算結果の
6/// オーバーフローは検査しないため、各型に収まる入力を使ってください。
7///
8/// # Examples
9///
10/// ```
11/// use atcoder::math::fast_math::FastMath;
12///
13/// assert_eq!(3_u64.fast_pow(4), 81);
14/// assert_eq!(2_u64.mod_pow(10, 1_000), 24);
15/// assert_eq!(12_u64.gcd(18), 6);
16/// ```
17pub trait FastMath: Sized + Copy {
18    /// `self` の `exponent` 乗を返します。
19    ///
20    /// `exponent` が正の場合の計算量は `O(log exponent)` です。指数が `0` 以下なら
21    /// `1` を返します（符号付き型の負指数も `1` になります）。積のオーバーフローは検査しません。
22    fn fast_pow(self, exponent: Self) -> Self;
23
24    /// 二分累乗法で `self.pow(exponent) mod modulus` を計算します。
25    ///
26    /// 非負の `exponent` と正の `modulus` を指定してください。計算量は `O(log exponent)`、
27    /// `modulus == 0` の場合は panic します。剰余前の積は通常の整数乗算なので、型の範囲を
28    /// 超えるとオーバーフローします。符号付き型では負数の剰余は非負に正規化されません。
29    fn mod_pow(self, exponent: Self, modulus: Self) -> Self;
30
31    /// `modulus` における乗法逆元を返します。
32    ///
33    /// `modulus >= 2` が素数で、`self` がその法で `0` でない場合に限り、
34    /// フェルマーの小定理に基づく逆元を返します。計算量は `O(log modulus)` です。
35    /// この条件を満たさない入力の結果は逆元として保証されません。
36    fn mod_inv(self, modulus: Self) -> Self;
37
38    /// ユークリッドの互除法で最大公約数を返します。通常の非負整数について計算量は
39    /// `O(log min(self, rhs))` です。符号付き型では負数を渡さないでください。
40    fn gcd(self, rhs: Self) -> Self;
41
42    /// `self / gcd(self, rhs) * rhs` により最小公倍数を返します。いずれかが `0` なら `0` です。
43    /// 乗算のオーバーフローは検査しません。符号付き型では非負の値を渡してください。
44    fn lcm(self, rhs: Self) -> Self;
45}
46
47macro_rules! impl_fast_math {
48    ($($integer:ty),* $(,)?) => {
49        $(
50            impl FastMath for $integer {
51                fn fast_pow(mut self, mut exponent: Self) -> Self {
52                    let mut result = 1;
53                    while exponent > 0 {
54                        if exponent & 1 == 1 {
55                            result *= self;
56                        }
57                        self *= self;
58                        exponent >>= 1;
59                    }
60                    result
61                }
62
63                fn mod_pow(mut self, mut exponent: Self, modulus: Self) -> Self {
64                    self %= modulus;
65                    let mut result = 1 % modulus;
66                    while exponent > 0 {
67                        if exponent & 1 == 1 {
68                            result = result * self % modulus;
69                        }
70                        self = self * self % modulus;
71                        exponent >>= 1;
72                    }
73                    result
74                }
75
76                fn mod_inv(self, modulus: Self) -> Self {
77                    self.mod_pow(modulus - 2, modulus)
78                }
79
80                fn gcd(self, rhs: Self) -> Self {
81                    let mut a = self;
82                    let mut b = rhs;
83                    while b != 0 {
84                        let remainder = a % b;
85                        a = b;
86                        b = remainder;
87                    }
88                    a
89                }
90
91                fn lcm(self, rhs: Self) -> Self {
92                    if self == 0 || rhs == 0 {
93                        0
94                    } else {
95                        self / self.gcd(rhs) * rhs
96                    }
97                }
98            }
99        )*
100    };
101}
102
103impl_fast_math!(i32, i64, isize, u32, u64, usize);
104
105#[cfg(test)]
106mod tests {
107    use super::FastMath;
108
109    #[test]
110    fn computes_integer_operations() {
111        assert_eq!(2_usize.fast_pow(10), 1024);
112        assert_eq!(3_usize.mod_inv(7), 5);
113        assert_eq!(21_usize.gcd(14), 7);
114        assert_eq!(6_usize.lcm(15), 30);
115    }
116}
atcoder/lib/math/fast_math.rs

atcoder/lib/math/
fast_math.rs
