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}