competitive/math/
mint_matrix.rs

1use super::{
2    AddMulOperation, DotProduct, MInt, MIntConvert, MIntDotProduct, Matrix, MemorizedFactorial,
3    One, Xorshift, Zero,
4};
5
6pub trait MIntMatrix<M>
7where
8    M: MIntDotProduct,
9{
10    /// det(self + other * x)
11    fn determinant_linear(self, other: Self) -> Option<Vec<MInt<M>>>
12    where
13        M: MIntConvert<usize> + MIntConvert<u64>;
14
15    fn pow_frobenius(self, k: usize) -> Self
16    where
17        M: MIntConvert<u64>;
18}
19
20impl<M> MIntMatrix<M> for Matrix<AddMulOperation<MInt<M>>>
21where
22    M: MIntDotProduct,
23{
24    fn determinant_linear(mut self, other: Self) -> Option<Vec<MInt<M>>>
25    where
26        M: MIntConvert<usize> + MIntConvert<u64>,
27    {
28        let mut rng = Xorshift::new();
29        let a = MInt::from(rng.rand64());
30        let n = self.data.len();
31        for i in 0..n {
32            for j in 0..n {
33                self[i][j] += other[i][j] * a;
34            }
35        }
36        let mut f = other.determinant_linear_non_singular(self)?;
37        f.reverse();
38        Some(taylor_shift::<M>(f, -a))
39    }
40
41    fn pow_frobenius(self, k: usize) -> Self
42    where
43        M: MIntConvert<u64>,
44    {
45        assert_eq!(self.shape.0, self.shape.1);
46        let a = self.transpose();
47        let mut rng = Xorshift::new();
48        let f = loop {
49            if let Some(f) = frobenius_decomposition(&a, &mut rng) {
50                break f;
51            }
52        };
53        let fk = f.pow(k);
54        let n = f.t.shape.0;
55        if f.blocks
56            .iter()
57            .map(|p| (p.0.len() - 1).pow(2))
58            .sum::<usize>()
59            * 4
60            <= n * n
61        {
62            let mut ft = Matrix::zeros((n, n));
63            let mut first = 0;
64            for p in &f.blocks {
65                let d = p.0.len() - 1;
66                for i in first..first + d {
67                    for j in first..first + d {
68                        MInt::add_scaled_assign(&mut ft[i], &f.t[j], &fk[i][j]);
69                    }
70                }
71                first += d;
72            }
73            &f.t_inv * &ft
74        } else {
75            &(&f.t_inv * &fk) * &f.t
76        }
77    }
78}
79
80impl<M> Matrix<AddMulOperation<MInt<M>>>
81where
82    M: MIntDotProduct,
83{
84    fn determinant_linear_non_singular(mut self, mut other: Self) -> Option<Vec<MInt<M>>>
85    where
86        M: MIntDotProduct,
87    {
88        let n = self.data.len();
89        let mut f = MInt::one();
90        for d in 0..n {
91            let i = other.data.iter().position(|other| !other[d].is_zero())?;
92            if i != d {
93                self.data.swap(i, d);
94                other.data.swap(i, d);
95                f = -f;
96            }
97            f *= other[d][d];
98            let r = other[d][d].inv();
99            for j in 0..n {
100                self[d][j] *= r;
101                other[d][j] *= r;
102            }
103            assert!(other[d][d].is_one());
104            for i in d + 1..n {
105                let a = other[i][d];
106                for k in 0..n {
107                    self[i][k] = self[i][k] - a * self[d][k];
108                    other[i][k] = other[i][k] - a * other[d][k];
109                }
110            }
111            for j in d + 1..n {
112                let a = other[d][j];
113                for k in 0..n {
114                    self[k][j] = self[k][j] - a * self[k][d];
115                    other[k][j] = other[k][j] - a * other[k][d];
116                }
117            }
118        }
119        for s in self.data.iter_mut() {
120            for s in s.iter_mut() {
121                *s = -*s;
122            }
123        }
124        let mut p = self.characteristic_polynomial();
125        for p in p.iter_mut() {
126            *p *= f;
127        }
128        Some(p)
129    }
130}
131
132struct EchelonRow<M>
133where
134    M: MIntDotProduct,
135{
136    pivot: usize,
137    inv: MInt<M>,
138    row: Vec<MInt<M>>,
139}
140
141struct Polynomial<M>(Vec<MInt<M>>)
142where
143    M: MIntDotProduct;
144
145struct FrobeniusDecomposition<M>
146where
147    M: MIntDotProduct,
148{
149    t: Matrix<AddMulOperation<MInt<M>>>,
150    t_inv: Matrix<AddMulOperation<MInt<M>>>,
151    blocks: Vec<Polynomial<M>>,
152}
153
154impl<M> EchelonRow<M>
155where
156    M: MIntDotProduct,
157{
158    fn reduce(&self, row: &mut [MInt<M>]) {
159        let a = -row[self.pivot] * self.inv;
160        if a.is_zero() {
161            return;
162        }
163        let end = self.row.len();
164        MInt::add_scaled_assign(&mut row[self.pivot..end], &self.row[self.pivot..], &a);
165    }
166}
167
168fn generate_frobenius_block<M>(
169    a: &Matrix<AddMulOperation<MInt<M>>>,
170    mut v: Vec<MInt<M>>,
171    rows: &mut Vec<EchelonRow<M>>,
172    t: &mut Vec<Vec<MInt<M>>>,
173) -> Polynomial<M>
174where
175    M: MIntDotProduct,
176{
177    let n = a.shape.0;
178    loop {
179        let mut row = vec![MInt::zero(); n + rows.len() + 1];
180        let (x, c) = row.split_at_mut(n);
181        x.copy_from_slice(&v);
182        c[rows.len()] = MInt::one();
183        for r in rows.iter() {
184            r.reduce(&mut row);
185        }
186        if let Some(pivot) = row[..n].iter().position(|x| !x.is_zero()) {
187            t.push(v);
188            let u = t.last().unwrap();
189            v = a.data.iter().map(|row| MInt::dot_product(u, row)).collect();
190            rows.push(EchelonRow {
191                pivot,
192                inv: row[pivot].inv(),
193                row,
194            });
195        } else {
196            let mut p = row.split_off(n);
197            while p.last().is_some_and(|x| x.is_zero()) {
198                p.pop();
199            }
200            return Polynomial(p);
201        }
202    }
203}
204
205impl<M> Polynomial<M>
206where
207    M: MIntDotProduct,
208{
209    fn exact_div(mut self, rhs: &Self) -> Option<Self> {
210        let mut q = vec![MInt::zero(); self.0.len() - rhs.0.len() + 1];
211        let inv = rhs.0.last().unwrap().inv();
212        for i in (0..q.len()).rev() {
213            q[i] = self.0[i + rhs.0.len() - 1] * inv;
214            MInt::add_scaled_assign(&mut self.0[i..i + rhs.0.len()], &rhs.0, &-q[i]);
215        }
216        self.0.iter().all(|x| x.is_zero()).then_some(Self(q))
217    }
218
219    fn square_mod(&self, p: &Self) -> Self {
220        let d = p.0.len() - 1;
221        let mut c = vec![MInt::zero(); 2 * d - 1];
222        for (i, &x) in self.0.iter().enumerate() {
223            MInt::add_scaled_assign(&mut c[i..2 * i], &self.0[..i], &(x + x));
224            c[2 * i] += x * x;
225        }
226        for i in (d..c.len()).rev() {
227            let x = c[i];
228            MInt::add_scaled_assign(&mut c[i - d..=i], &p.0, &-x);
229        }
230        c.truncate(d);
231        Self(c)
232    }
233
234    fn x_pow_mod(&self, k: usize) -> Self {
235        let d = self.0.len() - 1;
236        if d == 1 {
237            return Self(vec![(-self.0[0]).pow(k)]);
238        }
239        let mut r = Self(vec![MInt::zero(); d]);
240        r.0[0] = MInt::one();
241        for bit in (0..usize::BITS - k.leading_zeros()).rev() {
242            r = r.square_mod(self);
243            if k >> bit & 1 != 0 {
244                let x = r.0[d - 1];
245                for i in (1..d).rev() {
246                    r.0[i] = r.0[i - 1] - x * self.0[i];
247                }
248                r.0[0] = -x * self.0[0];
249            }
250        }
251        r
252    }
253}
254
255fn frobenius_decomposition<M>(
256    a: &Matrix<AddMulOperation<MInt<M>>>,
257    rng: &mut Xorshift,
258) -> Option<FrobeniusDecomposition<M>>
259where
260    M: MIntDotProduct + MIntConvert<u64>,
261{
262    let n = a.shape.0;
263    let mut rows = Vec::with_capacity(n);
264    let mut t = Vec::with_capacity(n);
265    let mut blocks: Vec<Polynomial<M>> = Vec::new();
266    while rows.len() < n {
267        let s = rows.len();
268        let v = (0..n).map(|_| MInt::from(rng.rand64())).collect();
269        let c = generate_frobenius_block(a, v, &mut rows, &mut t);
270        if rows.len() == s {
271            continue;
272        }
273        let p = Polynomial(c.0[s..].to_vec());
274        if c.0[..s].iter().any(|x| !x.is_zero()) {
275            let q = c.exact_div(&p)?;
276            let d = rows.len() - s;
277            let mut coefficients = q.0[..s].to_vec();
278            let mut shifts = Vec::with_capacity(d);
279            for _ in 0..d {
280                shifts.push(coefficients.clone());
281                let mut first = 0;
282                for block in &blocks {
283                    let len = block.0.len() - 1;
284                    let c = &mut coefficients[first..first + len];
285                    let last = c[len - 1];
286                    for j in (1..len).rev() {
287                        c[j] = c[j - 1] - last * block.0[j];
288                    }
289                    c[0] = -last * block.0[0];
290                    first += len;
291                }
292            }
293            let shifts: Matrix<AddMulOperation<MInt<M>>> = Matrix::from_vec(shifts);
294            if d < 32 {
295                let (previous, current) = t.split_at_mut(s);
296                for (shift, row) in shifts.data.iter().zip(current) {
297                    for (factor, source) in shift.iter().zip(previous.iter()) {
298                        if !factor.is_zero() {
299                            MInt::add_scaled_assign(row, source, factor);
300                        }
301                    }
302                }
303            } else {
304                let previous = Matrix::from_vec(t[..s].to_vec());
305                let correction = &shifts * &previous;
306                for (row, correction) in t[s..].iter_mut().zip(&correction.data) {
307                    for (x, &y) in row.iter_mut().zip(correction) {
308                        *x += y;
309                    }
310                }
311            }
312            for row in &mut rows[s..] {
313                // Keep the reduced vector fixed: T_new += S*T_old gives C_old -= C_new*S.
314                let (previous, current) = row.row[n..].split_at_mut(s);
315                for (&x, shift) in current.iter().zip(&shifts.data) {
316                    MInt::add_scaled_assign(previous, shift, &-x);
317                }
318            }
319        }
320        blocks.push(p);
321    }
322
323    let mut t_inv = vec![vec![MInt::zero(); n]; n];
324    for i in (0..n).rev() {
325        let row = &rows[i];
326        let mut c = row.row[n..].to_vec();
327        c.resize(n, MInt::zero());
328        for x in &mut c {
329            *x *= row.inv;
330        }
331        for next in &rows[i + 1..] {
332            let factor = -row.row[next.pivot] * row.inv;
333            if !factor.is_zero() {
334                MInt::add_scaled_assign(&mut c, &t_inv[next.pivot], &factor);
335            }
336        }
337        t_inv[row.pivot] = c;
338    }
339    Some(FrobeniusDecomposition {
340        t: Matrix::from_vec(t),
341        t_inv: Matrix::from_vec(t_inv),
342        blocks,
343    })
344}
345
346impl<M> FrobeniusDecomposition<M>
347where
348    M: MIntDotProduct,
349{
350    fn pow(&self, k: usize) -> Matrix<AddMulOperation<MInt<M>>> {
351        let n = self.t.shape.0;
352        let mut a = vec![vec![MInt::zero(); n]; n];
353        let mut s = 0;
354        for p in &self.blocks {
355            let d = p.0.len() - 1;
356            let mut c = p.x_pow_mod(k).0;
357            for row in &mut a[s..s + d] {
358                row[s..s + d].copy_from_slice(&c);
359                let x = c[d - 1];
360                for i in (1..d).rev() {
361                    c[i] = c[i - 1] - x * p.0[i];
362                }
363                c[0] = -x * p.0[0];
364            }
365            s += d;
366        }
367        Matrix::from_vec(a)
368    }
369}
370
371fn taylor_shift<M>(f: Vec<MInt<M>>, a: MInt<M>) -> Vec<MInt<M>>
372where
373    M: MIntConvert<usize>,
374{
375    let n = f.len();
376    if n == 0 {
377        return f;
378    }
379    let mf = MemorizedFactorial::new(n);
380    let mut res = vec![MInt::<M>::zero(); n];
381    let mut apow = vec![MInt::<M>::one(); n];
382    for i in 1..n {
383        apow[i] = apow[i - 1] * a;
384    }
385    for j in 0..n {
386        if f[j].is_zero() {
387            continue;
388        }
389        for k in 0..=j {
390            res[k] += f[j] * apow[j - k] * mf.combination(j, k);
391        }
392    }
393    res
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399    use crate::{math::lagrange_interpolation_polynomial, num::montgomery::MInt998244353, rand};
400
401    #[test]
402    fn test_determinant_linear() {
403        let mut rng = Xorshift::default();
404        for _ in 0..100 {
405            rand!(rng, n: 1..30, m0: [[0..998244353; n]; n], m1: [[0..998244353; n]; n]);
406            let m0 = Matrix::<AddMulOperation<_>>::from_vec(m0)
407                .map::<AddMulOperation<MInt998244353>, _>(|&x| MInt998244353::new(x));
408            let m1 = Matrix::<AddMulOperation<_>>::from_vec(m1)
409                .map::<AddMulOperation<MInt998244353>, _>(|&x| MInt998244353::new(x));
410            let f = m0.clone().determinant_linear(m1.clone()).unwrap();
411
412            let d: Vec<_> = (0..=n)
413                .map(|k| {
414                    let mut mat = Matrix::<AddMulOperation<_>>::new_with((n, n), |i, j| {
415                        m0[i][j] + m1[i][j] * MInt998244353::from(k)
416                    });
417                    mat.determinant()
418                })
419                .collect();
420            let (x, y): (Vec<_>, Vec<_>) = (0..=n).map(|k| (MInt998244353::from(k), d[k])).unzip();
421            let g = lagrange_interpolation_polynomial(&x, &y);
422            assert_eq!(f, g);
423        }
424    }
425
426    #[test]
427    fn test_pow_frobenius() {
428        let mut rng = Xorshift::default();
429        for iteration in 0..100 {
430            let n = if iteration < 16 {
431                rng.random(32..100)
432            } else {
433                rng.random(0..30)
434            };
435            let k = rng.random(0..1_000_000_000);
436            rand!(rng, data: [[0..998244353; n]; n]);
437            let matrix = Matrix::<AddMulOperation<_>>::from_vec(data)
438                .map::<AddMulOperation<MInt998244353>, _>(|&x| MInt998244353::new(x));
439            assert_eq!(matrix.clone().pow(k), matrix.pow_frobenius(k));
440
441            let scalar: MInt998244353 = rng.random(..);
442            let matrix: Matrix<AddMulOperation<MInt998244353>> =
443                Matrix::new_with((n, n), |i, j| if i == j { scalar } else { MInt::zero() });
444            assert_eq!(matrix.clone().pow(k), matrix.pow_frobenius(k));
445
446            let mut matrix: Matrix<AddMulOperation<MInt998244353>> =
447                Matrix::new_with((n, n), |i, j| {
448                    if i == j {
449                        scalar
450                    } else if i + 1 == j && rng.gen_bool(0.8) {
451                        MInt::one()
452                    } else {
453                        MInt::zero()
454                    }
455                });
456            if n >= 2 {
457                for _ in 0..4 * n {
458                    let i = rng.random(..n);
459                    let j = (i + rng.random(1..n)) % n;
460                    let factor: MInt998244353 = rng.random(..);
461                    for k in 0..n {
462                        let x = factor * matrix[j][k];
463                        matrix[i][k] += x;
464                    }
465                    for row in &mut matrix.data {
466                        let x = factor * row[i];
467                        row[j] -= x;
468                    }
469                }
470            }
471            assert_eq!(matrix.clone().pow(k), matrix.pow_frobenius(k));
472        }
473    }
474}
competitive/math/mint_matrix.rs

competitive/math/
mint_matrix.rs
