competitive/math/
linear_congruence.rs

1use super::Unsigned;
2
3/// return: (y,z)
4///
5/// ax = b mod m, where x = y mod z
6pub fn solve_linear_congruence<T>(a: T, b: T, m: T) -> Option<(T, T)>
7where
8    T: Unsigned,
9{
10    let g = a.gcd(m);
11    if b % g != T::zero() {
12        return None;
13    }
14    let (a, b, m) = (a / g, b / g, m / g);
15    Some((b * a.modinv(m) % m, m))
16}
17
18/// return: (y,z)
19///
20/// forall (a,b,m), ax = b mod m, where x = y mod z
21pub fn solve_simultaneous_linear_congruence<T, I>(abm: I) -> Option<(T, T)>
22where
23    T: Unsigned,
24    I: IntoIterator<Item = (T, T, T)>,
25{
26    let mut x = T::zero();
27    let mut m0 = T::one();
28    for (a, b, m) in abm {
29        let mut b = b + m - a * x % m;
30        if b >= m {
31            b -= m;
32        }
33        let a = a * m0;
34        let (y, z) = solve_linear_congruence(a, b, m)?;
35        x += y * m0;
36        m0 *= z;
37    }
38    x %= m0;
39    Some((x, m0))
40}
41
42#[test]
43fn test_linear_congruence() {
44    use crate::math::lcm;
45    use crate::tools::Xorshift;
46    const N: usize = 5;
47    const Q: usize = 1_000;
48    let mut rng = Xorshift::new();
49    for _ in 0..Q {
50        let abm: Vec<_> = (0..N)
51            .map(|_| {
52                let m = rng.random(2u64..=20);
53                (rng.random(1..m), rng.random(0..m), m)
54            })
55            .collect();
56        if let Some((x, m0)) = solve_simultaneous_linear_congruence(abm.iter().cloned()) {
57            assert!(x < m0);
58            for (a, b, m) in abm.iter().cloned() {
59                assert!(a * x % m == b);
60            }
61        } else {
62            let m0 = abm[1..].iter().fold(abm[0].2, |x, y| lcm(x, y.2));
63            let x = (0..m0).find(|&x| abm.iter().cloned().all(|(a, b, m)| a * x % m == b));
64            assert_eq!(x, None);
65        }
66    }
67}
competitive/math/linear_congruence.rs

competitive/math/
linear_congruence.rs
