competitive/math/
mint_matrix.rs

1use super::{
2    AddMulOperation, MInt, MIntBase, MIntConvert, Matrix, MemorizedFactorial, One, Xorshift, Zero,
3};
4
5pub trait MIntMatrix<M>
6where
7    M: MIntBase,
8{
9    /// det(self + other * x)
10    fn determinant_linear(self, other: Self) -> Option<Vec<MInt<M>>>
11    where
12        M: MIntConvert<usize> + MIntConvert<u64>;
13}
14
15impl<M> MIntMatrix<M> for Matrix<AddMulOperation<MInt<M>>>
16where
17    M: MIntBase,
18{
19    fn determinant_linear(mut self, other: Self) -> Option<Vec<MInt<M>>>
20    where
21        M: MIntConvert<usize> + MIntConvert<u64>,
22    {
23        let mut rng = Xorshift::new();
24        let a = MInt::from(rng.rand64());
25        let n = self.data.len();
26        for i in 0..n {
27            for j in 0..n {
28                self[i][j] += other[i][j] * a;
29            }
30        }
31        let mut f = other.determinant_linear_non_singular(self)?;
32        f.reverse();
33        Some(taylor_shift::<M>(f, -a))
34    }
35}
36
37impl<M> Matrix<AddMulOperation<MInt<M>>>
38where
39    M: MIntBase,
40{
41    fn determinant_linear_non_singular(mut self, mut other: Self) -> Option<Vec<MInt<M>>>
42    where
43        M: MIntBase,
44    {
45        let n = self.data.len();
46        let mut f = MInt::one();
47        for d in 0..n {
48            let i = other.data.iter().position(|other| !other[d].is_zero())?;
49            if i != d {
50                self.data.swap(i, d);
51                other.data.swap(i, d);
52                f = -f;
53            }
54            f *= other[d][d];
55            let r = other[d][d].inv();
56            for j in 0..n {
57                self[d][j] *= r;
58                other[d][j] *= r;
59            }
60            assert!(other[d][d].is_one());
61            for i in d + 1..n {
62                let a = other[i][d];
63                for k in 0..n {
64                    self[i][k] = self[i][k] - a * self[d][k];
65                    other[i][k] = other[i][k] - a * other[d][k];
66                }
67            }
68            for j in d + 1..n {
69                let a = other[d][j];
70                for k in 0..n {
71                    self[k][j] = self[k][j] - a * self[k][d];
72                    other[k][j] = other[k][j] - a * other[k][d];
73                }
74            }
75        }
76        for s in self.data.iter_mut() {
77            for s in s.iter_mut() {
78                *s = -*s;
79            }
80        }
81        let mut p = self.characteristic_polynomial();
82        for p in p.iter_mut() {
83            *p *= f;
84        }
85        Some(p)
86    }
87}
88
89fn taylor_shift<M>(f: Vec<MInt<M>>, a: MInt<M>) -> Vec<MInt<M>>
90where
91    M: MIntConvert<usize>,
92{
93    let n = f.len();
94    if n == 0 {
95        return f;
96    }
97    let mf = MemorizedFactorial::new(n);
98    let mut res = vec![MInt::<M>::zero(); n];
99    let mut apow = vec![MInt::<M>::one(); n];
100    for i in 1..n {
101        apow[i] = apow[i - 1] * a;
102    }
103    for j in 0..n {
104        if f[j].is_zero() {
105            continue;
106        }
107        for k in 0..=j {
108            res[k] += f[j] * apow[j - k] * mf.combination(j, k);
109        }
110    }
111    res
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use crate::{math::lagrange_interpolation_polynomial, num::montgomery::MInt998244353, rand};
118
119    #[test]
120    fn test_determinant_linear() {
121        let mut rng = Xorshift::default();
122        for _ in 0..100 {
123            rand!(rng, n: 1..30, m0: [[0..998244353; n]; n], m1: [[0..998244353; n]; n]);
124            let m0 = Matrix::<AddMulOperation<_>>::from_vec(m0)
125                .map::<AddMulOperation<MInt998244353>, _>(|&x| MInt998244353::new(x));
126            let m1 = Matrix::<AddMulOperation<_>>::from_vec(m1)
127                .map::<AddMulOperation<MInt998244353>, _>(|&x| MInt998244353::new(x));
128            let f = m0.clone().determinant_linear(m1.clone()).unwrap();
129
130            let d: Vec<_> = (0..=n)
131                .map(|k| {
132                    let mut mat = Matrix::<AddMulOperation<_>>::new_with((n, n), |i, j| {
133                        m0[i][j] + m1[i][j] * MInt998244353::from(k)
134                    });
135                    mat.determinant()
136                })
137                .collect();
138            let (x, y): (Vec<_>, Vec<_>) = (0..=n).map(|k| (MInt998244353::from(k), d[k])).unzip();
139            let g = lagrange_interpolation_polynomial(&x, &y);
140            assert_eq!(f, g);
141        }
142    }
143}