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#[cfg(test)]
40mod tests {
41    use super::*;
42    use crate::{math::lcm, tools::Xorshift};
43
44    #[test]
45    fn test_linear_congruence() {
46        const N: usize = 5;
47        const Q: usize = 1_000;
48        let mut rng = Xorshift::default();
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    }
68}
competitive/math/linear_congruence.rs

competitive/math/
linear_congruence.rs
