Skip to main content

nekolib/math/
mod_pow.rs

1//! 冪乗。
2
3/// 冪乗。
4///
5/// $a^b \\bmod n$ を返す。
6///
7/// # Complexity
8/// $O(\\log(b))$ time.
9///
10/// # Examples
11/// ```
12/// use nekolib::math::ModPow;
13///
14/// assert_eq!(3_u64.mod_pow(14, 10), 9);
15/// assert_eq!(2_u64.mod_pow(11, 1024), 0);
16/// assert_eq!(0_u64.mod_pow(0, 1), 0);
17/// ```
18pub trait ModPow {
19    fn mod_pow(self, b: Self, n: Self) -> Self;
20}
21
22macro_rules! impl_uint {
23    ($t:ty) => {
24        impl ModPow for $t {
25            fn mod_pow(self, mut b: Self, n: Self) -> Self {
26                if n == 1 {
27                    return 0; // in case 0^0
28                }
29                let mut res = 1;
30                let mut a = self % n;
31                while b > 0 {
32                    if b & 1 == 1 {
33                        res = res * a % n;
34                    }
35                    a = a * a % n;
36                    b >>= 1;
37                }
38                res
39            }
40        }
41    };
42    ( $($t:ty)* ) => { $(impl_uint!($t);)* };
43}
44
45impl_uint!(u8 u16 u32 u64 u128 usize);
46
47#[test]
48fn test() {
49    for n in 1_u64..=30 {
50        for a in 0..30 {
51            let mut expected = 1 % n;
52            for b in 0..30 {
53                assert_eq!(a.mod_pow(b, n), expected);
54                expected = expected * a % n;
55            }
56        }
57    }
58}