aoc/util/math.rs
1//! Extended mathematical operations.
2//!
3//! * [Greatest common divisor](https://en.wikipedia.org/wiki/Greatest_common_divisor)
4//! of 2 numbers using the
5//! [Euclidean algorithm](https://en.wikipedia.org/wiki/Euclidean_algorithm).
6//!
7//! * [Least common multiple](https://en.wikipedia.org/wiki/Least_common_multiple)
8//!
9//! * [Modular exponentation](https://en.wikipedia.org/wiki/Modular_exponentiation).
10//! Calculates bᵉ mod m efficiently using
11//! [exponentiation by squaring](https://en.wikipedia.org/wiki/Exponentiation_by_squaring).
12//!
13//! * [Modular multiplicative inverse](https://en.wikipedia.org/wiki/Modular_multiplicative_inverse)
14//! calculated using the [extended Euclidean algorithm](https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm).
15//!
16//! * [Integer square root](https://en.wikipedia.org/wiki/Integer_square_root).
17use crate::util::integer::*;
18
19pub trait IntegerMathOps<T: Integer<T>> {
20 fn gcd(self, b: T) -> T;
21 fn lcm(self, b: T) -> T;
22 fn mod_pow(self, e: T, m: T) -> T;
23}
24
25pub trait SignedMathOps<T: Signed<T>> {
26 fn mod_inv(self, m: T) -> Option<T>;
27}
28
29impl<T: Integer<T>> IntegerMathOps<T> for T {
30 /// Greatest common divisor
31 #[inline]
32 fn gcd(self, mut b: T) -> T {
33 let mut a = self;
34
35 while b != T::ZERO {
36 (a, b) = (b, a % b);
37 }
38
39 a
40 }
41
42 /// Least common multiple
43 #[inline]
44 fn lcm(self, b: T) -> T {
45 self * (b / self.gcd(b))
46 }
47
48 /// Modular exponentiation
49 #[inline]
50 fn mod_pow(self, mut e: T, m: T) -> T {
51 let mut base = self;
52 let mut result = T::ONE;
53
54 while e > T::ZERO {
55 if e & T::ONE == T::ONE {
56 result = (result * base) % m;
57 }
58 base = (base * base) % m;
59 e = e >> 1;
60 }
61
62 result
63 }
64}
65
66impl<T: Signed<T>> SignedMathOps<T> for T {
67 /// Modular multiplicative inverse
68 #[inline]
69 fn mod_inv(self, m: T) -> Option<T> {
70 let mut t = T::ZERO;
71 let mut new_t = T::ONE;
72 let mut r = m;
73 let mut new_r = self;
74
75 while new_r != T::ZERO {
76 let quotient = r / new_r;
77 (t, new_t) = (new_t, t - quotient * new_t);
78 (r, new_r) = (new_r, r - quotient * new_r);
79 }
80
81 if r > T::ONE {
82 return None;
83 }
84 if t < T::ZERO {
85 t = t + m;
86 }
87
88 Some(t)
89 }
90}