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 exponentiation](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).
15use crate::util::integer::*;
16
17pub trait IntegerMathOps<T: Integer<T>> {
18    #[must_use]
19    fn gcd(self, b: T) -> T;
20    #[must_use]
21    fn lcm(self, b: T) -> T;
22    #[must_use]
23    fn mod_pow(self, e: T, m: T) -> T;
24}
25
26pub trait SignedMathOps<T: Signed<T>> {
27    #[must_use]
28    fn mod_inv(self, m: T) -> Option<T>;
29}
30
31impl<T: Integer<T>> IntegerMathOps<T> for T {
32    /// Greatest common divisor
33    #[inline]
34    fn gcd(self, mut b: T) -> T {
35        let mut a = self;
36
37        while b != T::ZERO {
38            (a, b) = (b, a % b);
39        }
40
41        a
42    }
43
44    /// Least common multiple
45    #[inline]
46    fn lcm(self, b: T) -> T {
47        self * (b / self.gcd(b))
48    }
49
50    /// Modular exponentiation
51    #[inline]
52    fn mod_pow(self, mut e: T, m: T) -> T {
53        let mut base = self;
54        let mut result = T::ONE;
55
56        while e > T::ZERO {
57            if e & T::ONE == T::ONE {
58                result = (result * base) % m;
59            }
60            base = (base * base) % m;
61            e = e >> 1;
62        }
63
64        result
65    }
66}
67
68impl<T: Signed<T>> SignedMathOps<T> for T {
69    /// Modular multiplicative inverse
70    #[inline]
71    fn mod_inv(self, m: T) -> Option<T> {
72        let mut t = T::ZERO;
73        let mut new_t = T::ONE;
74        let mut r = m;
75        let mut new_r = self;
76
77        while new_r != T::ZERO {
78            let quotient = r / new_r;
79            (t, new_t) = (new_t, t - quotient * new_t);
80            (r, new_r) = (new_r, r - quotient * new_r);
81        }
82
83        if r > T::ONE {
84            return None;
85        }
86        if t < T::ZERO {
87            t = t + m;
88        }
89
90        Some(t)
91    }
92}