haar_lib/linalg/mod_p/
lineq.rs

1//! $\mathbb{Z} / p \mathbb{Z}$($p$は素数)上の連立一次方程式
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/system_of_linear_equations>
5use crate::{linalg::mod_p::gaussian_elim::*, num::ff::*};
6
7/// $\mathbb{Z} / p \mathbb{Z}$($p$は素数)上で連立一次方程式$A \boldsymbol{x} = \boldsymbol{b}$を解く。
8///
9/// ここで、$A$は$n \times m$の行列、$\boldsymbol{x}$は$m$行の縦ベクトル、$\boldsymbol{b}$は$n$行の縦ベクトル。
10///
11/// 連立方程式が解をもたないとき、`None`を返す。
12/// そうでなければ、`Some((sol, bases))`を返す。
13///
14/// `sol`は$m$行のベクトル、`bases`は`dim`個の$m$行のベクトルで、
15/// 連立方程式の解は、`bases`の要素の線型結合と`sol`の和で表される。
16pub fn lineq<F>(
17    mut a: Vec<Vec<F::Element>>,
18    b: Vec<F::Element>,
19    modulo: &F,
20) -> Option<(Vec<F::Element>, Vec<Vec<F::Element>>)>
21where
22    F: FF,
23    F::Element: FFElem,
24{
25    let n = a.len();
26    assert_eq!(b.len(), n);
27
28    let Some(m) = a.first().map(|a| a.len()) else {
29        panic!("行数は0以上でなければならない。")
30    };
31    assert!(a.iter().all(|r| r.len() == m));
32
33    for (r, bi) in a.iter_mut().zip(b.iter()) {
34        r.push(*bi);
35    }
36
37    let (rank, mut a) = gaussian_elim(a);
38
39    let dim = m - rank;
40
41    let b: Vec<_> = a.iter_mut().map(|r| r.pop().unwrap()).collect();
42
43    if rank > 0 && a[rank - 1].iter().all(|x| x.value() == 0) {
44        return None;
45    }
46
47    let mut index_zero = vec![];
48    let mut index_one = vec![];
49    {
50        let mut k = 0;
51        for ai in a.iter().take(rank) {
52            for (j, aij) in ai.iter().enumerate().take(m).skip(k) {
53                if aij.value() != 0 {
54                    index_one.push(j);
55                    k = j + 1;
56                    break;
57                }
58                index_zero.push(j);
59            }
60        }
61        for j in k..m {
62            index_zero.push(j);
63        }
64    }
65
66    assert_eq!(index_zero.len(), dim);
67    assert_eq!(index_one.len(), rank);
68
69    let mut sol = vec![modulo.zero(); m];
70    for (i, x) in b.into_iter().take(rank).enumerate() {
71        sol[index_one[i]] = x;
72    }
73
74    let mut bases = vec![vec![modulo.zero(); m]; dim];
75    for i in 0..rank {
76        for (j, &k) in index_zero.iter().enumerate() {
77            bases[j][index_one[i]] = -a[i][k];
78        }
79    }
80
81    for i in 0..dim {
82        bases[i][index_zero[i]] = modulo.one();
83    }
84
85    Some((sol, bases))
86}
haar_lib/linalg/mod_p/lineq.rs

haar_lib/linalg/mod_p/
lineq.rs
