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 fn gcd(self, mut b: T) -> T {
32 let mut a = self;
33
34 while b != T::ZERO {
35 (a, b) = (b, a % b);
36 }
37
38 a
39 }
40
41 // Least common multiple
42 fn lcm(self, b: T) -> T {
43 self * (b / self.gcd(b))
44 }
45
46 // Modular exponentation
47 fn mod_pow(self, mut e: T, m: T) -> T {
48 let mut b = self;
49 let mut c = T::ONE;
50
51 while e > T::ZERO {
52 if e & T::ONE == T::ONE {
53 c = (c * b) % m;
54 }
55 b = (b * b) % m;
56 e = e >> T::ONE;
57 }
58
59 c
60 }
61}
62
63impl<T: Signed<T>> SignedMathOps<T> for T {
64 // Modular multiplicative inverse
65 fn mod_inv(self, m: T) -> Option<T> {
66 let mut t = T::ZERO;
67 let mut newt = T::ONE;
68 let mut r = m;
69 let mut newr = self;
70
71 while newr != T::ZERO {
72 let quotient = r / newr;
73 (t, newt) = (newt, t - quotient * newt);
74 (r, newr) = (newr, r - quotient * newr);
75 }
76
77 if r > T::ONE {
78 return None;
79 }
80 if t < T::ZERO {
81 t = t + m;
82 }
83
84 Some(t)
85 }
86}