haar_lib/linalg/mod_p/
inverse.rs

1//! $\mathbb{Z} / p \mathbb{Z}$($p$は素数)上の逆行列
2use crate::num::ff::*;
3
4/// $\mathbb{Z} / p \mathbb{Z}$($p$は素数)上での逆行列を求める。
5///
6/// **Time complexity** $O(n^3)$
7pub fn inverse<F>(mut b: Vec<Vec<F::Element>>, modulo: &F) -> Option<Vec<Vec<F::Element>>>
8where
9    F: FF,
10    F::Element: FFElem,
11{
12    let n = b.len();
13
14    assert!(b.iter().all(|r| r.len() == n));
15
16    for (i, bi) in b.iter_mut().enumerate() {
17        bi.resize(2 * n, modulo.zero());
18        bi[i + n] = modulo.one();
19    }
20
21    for i in 0..n {
22        let q = (i..n).find(|&j| b[j][i].value() != 0)?;
23
24        b.swap(i, q);
25
26        let d = b[i][i].inv();
27
28        for x in b[i].iter_mut() {
29            *x *= d;
30        }
31
32        let d = b[i][i].inv();
33
34        let bi = b.swap_remove(i);
35
36        for bj in b.iter_mut() {
37            let d = bj[i] * d;
38
39            for (x, y) in bj.iter_mut().zip(bi.iter()) {
40                *x -= *y * d;
41            }
42        }
43
44        b.push(bi);
45        b.swap(i, n - 1);
46    }
47
48    let ret = b.into_iter().map(|a| a[n..].to_vec()).collect();
49
50    Some(ret)
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56    use crate::num::modint::*;
57
58    fn check(a: Vec<Vec<i64>>, m: u32, ans: Option<Vec<Vec<i64>>>) {
59        let m = ModIntBuilder::new(m);
60        let a = a
61            .into_iter()
62            .map(|b| b.into_iter().map(|x| m.from_i64(x)).collect())
63            .collect();
64        let ans = ans.map(|a| {
65            a.into_iter()
66                .map(|b| b.into_iter().map(|x| m.from_i64(x)).collect())
67                .collect()
68        });
69        assert_eq!(inverse(a, &m), ans);
70    }
71
72    #[test]
73    fn test() {
74        check(
75            vec![vec![3, 1, 4], vec![1, 5, 9], vec![2, 6, 5]],
76            998244353,
77            Some(vec![
78                vec![188557267, 255106890, 587855008],
79                vec![122007643, 987152749, 321656514],
80                vec![576763404, 310564910, 976061145],
81            ]),
82        );
83        check(
84            vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
85            998244353,
86            None,
87        );
88        check(
89            vec![vec![0, 1], vec![1, 0]],
90            998244353,
91            Some(vec![vec![0, 1], vec![1, 0]]),
92        );
93    }
94}
haar_lib/linalg/mod_p/inverse.rs

haar_lib/linalg/mod_p/
inverse.rs
