haar_lib/math/
crt.rs

1//! 中国剰余定理
2
3use crate::math::ext_gcd::ext_gcd;
4
5/// 二元の中国剰余定理
6///
7/// $$\begin{aligned}
8/// x \equiv b_1 \pmod {m_1} \\\\
9/// x \equiv b_2 \pmod {m_2}
10/// \end{aligned}$$
11/// を満たす$x \pmod {\mathrm{lcm}(m_1, m_2)}$が存在すれば、$x$と$\mathrm{lcm}(m_1, m_2)$を返す。
12/// そうでなければ、`None`を返す。
13pub fn crt((b1, m1): (i64, u64), (b2, m2): (i64, u64)) -> Option<(i64, u64)> {
14    let (d, p, _) = ext_gcd(m1, m2);
15
16    let m1 = m1 as i64;
17    let m2 = m2 as i64;
18
19    if (b2 - b1) % d != 0 {
20        return None;
21    }
22
23    let m = m1 / d * m2;
24    let t = ((b2 - b1) * p / d) % (m2 / d);
25    let r = (b1 + m1 * t + m) % m;
26
27    Some((r, m as u64))
28}
29
30/// 多元の中国剰余定理
31pub fn crt_vec(params: &[(i64, u64)]) -> Option<(i64, u64)> {
32    let mut _r = 0;
33    let mut _m = 1;
34
35    for &p in params {
36        match crt((_r, _m), p) {
37            Some((r, m)) => {
38                _r = r;
39                _m = m;
40            }
41            _ => return None,
42        }
43    }
44
45    Some((_r, _m))
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn test() {
54        assert_eq!(crt((2, 3), (3, 5)), Some((8, 15)));
55
56        // https://yukicoder.me/problems/447
57        assert_eq!(crt_vec(&[(10, 20), (10, 30), (30, 40)]), Some((70, 120)));
58        assert_eq!(crt_vec(&[(1, 2), (0, 4), (5, 17)]), None);
59    }
60}