Skip to main content

competitive/math/
black_box_mint_matrix.rs

1use super::{
2    AddMulOperation, BlackBoxMatrix, BlackBoxMatrixImpl, ConvolveSteps, DotProduct,
3    FormalPowerSeries, Fps, MInt, MIntConvert, MIntDotProduct, One, Xorshift, Zero,
4};
5
6pub trait BlackBoxMIntMatrix<M>: BlackBoxMatrix<AddMulOperation<MInt<M>>>
7where
8    M: MIntDotProduct<Inner = u32>
9        + MIntConvert<u32>
10        + MIntConvert<u64>
11        + MIntConvert<usize>
12        + MIntConvert<isize>,
13{
14    fn minimal_polynomial(&self) -> Vec<MInt<M>> {
15        assert_eq!(self.shape().0, self.shape().1);
16        let n = self.shape().0;
17        let mut rng = Xorshift::new();
18        let b: Vec<MInt<M>> = (0..n).map(|_| MInt::from(rng.rand64())).collect();
19        let u: Vec<MInt<M>> = (0..n).map(|_| MInt::from(rng.rand64())).collect();
20        let a: Vec<MInt<M>> = (0..2 * n)
21            .scan(b, |b, _| {
22                let a = MInt::dot_product(b, &u);
23                *b = self.apply(b);
24                Some(a)
25            })
26            .collect();
27        let polynomial: Fps<M> = FormalPowerSeries::berlekamp_massey(&a);
28        let mut p = polynomial.data;
29        p.reverse();
30        p
31    }
32
33    fn apply_pow<C>(&self, mut b: Vec<MInt<M>>, k: usize) -> Vec<MInt<M>>
34    where
35        C: ConvolveSteps<T = Vec<MInt<M>>>,
36    {
37        assert_eq!(self.shape().0, self.shape().1);
38        assert_eq!(self.shape().1, b.len());
39        let n = self.shape().0;
40        let p = self.minimal_polynomial();
41        let polynomial: FormalPowerSeries<MInt<M>, C> = FormalPowerSeries::from_vec(p);
42        let f = polynomial.pow_mod(k);
43        let mut res = vec![MInt::zero(); n];
44        for f in f {
45            for j in 0..n {
46                res[j] += f * b[j];
47            }
48            b = self.apply(&b);
49        }
50        res
51    }
52
53    fn black_box_determinant(&self) -> MInt<M> {
54        assert_eq!(self.shape().0, self.shape().1);
55        let n = self.shape().0;
56        let mut rng = Xorshift::new();
57        let d: Vec<MInt<M>> = (0..n).map(|_| MInt::from(rng.rand64())).collect();
58        let det_d = d.iter().fold(MInt::one(), |s, x| s * x);
59        let ad: BlackBoxMatrixImpl<AddMulOperation<MInt<M>>, _> =
60            BlackBoxMatrixImpl::new(self.shape(), |v: &[MInt<M>]| {
61                let mut w = self.apply(v);
62                for (w, d) in w.iter_mut().zip(&d) {
63                    *w *= d;
64                }
65                w
66            });
67        let p = ad.minimal_polynomial();
68        let det_ad = if n % 2 == 0 { p[0] } else { -p[0] };
69        det_ad / det_d
70    }
71
72    fn black_box_linear_equation(&self, mut b: Vec<MInt<M>>) -> Option<Vec<MInt<M>>> {
73        assert_eq!(self.shape().0, self.shape().1);
74        assert_eq!(self.shape().1, b.len());
75        let n = self.shape().0;
76        let p = self.minimal_polynomial();
77        if p.is_empty() || p[0].is_zero() {
78            return None;
79        }
80        let p0_inv = p[0].inv();
81        let mut x = vec![MInt::zero(); n];
82        for p in p.into_iter().skip(1) {
83            let p = -p * p0_inv;
84            for i in 0..n {
85                x[i] += p * b[i];
86            }
87            b = self.apply(&b);
88        }
89        Some(x)
90    }
91}
92
93impl<M, B> BlackBoxMIntMatrix<M> for B
94where
95    M: MIntDotProduct<Inner = u32>
96        + MIntConvert<u32>
97        + MIntConvert<u64>
98        + MIntConvert<usize>
99        + MIntConvert<isize>,
100    B: BlackBoxMatrix<AddMulOperation<MInt<M>>>,
101{
102}