haar_lib/linalg/mod_p/
lineq.rs

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

haar_lib/linalg/mod_p/
lineq.rs
