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}