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.mod_mul(a.mod_inv(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 b = b.mod_sub(a.mod_mul(x, m), m);
30        let a = a * m0;
31        let (y, z) = solve_linear_congruence(a, b, m)?;
32        x += y * m0;
33        m0 *= z;
34    }
35    x %= m0;
36    Some((x, m0))
37}
38
39#[test]
40fn test_linear_congruence() {
41    use crate::math::lcm;
42    use crate::tools::Xorshift;
43    const N: usize = 5;
44    const Q: usize = 1_000;
45    let mut rng = Xorshift::new();
46    for _ in 0..Q {
47        let abm: Vec<_> = (0..N)
48            .map(|_| {
49                let m = rng.random(2u64..=20);
50                (rng.random(1..m), rng.random(0..m), m)
51            })
52            .collect();
53        if let Some((x, m0)) = solve_simultaneous_linear_congruence(abm.iter().cloned()) {
54            assert!(x < m0);
55            for (a, b, m) in abm.iter().cloned() {
56                assert!(a * x % m == b);
57            }
58        } else {
59            let m0 = abm[1..].iter().fold(abm[0].2, |x, y| lcm(x, y.2));
60            let x = (0..m0).find(|&x| abm.iter().cloned().all(|(a, b, m)| a * x % m == b));
61            assert_eq!(x, None);
62        }
63    }
64}