pub trait SemiRing {
type T: Clone;
type Additive: AbelianMonoid<T = Self::T>;
type Multiplicative: Monoid<T = Self::T>;
// Provided methods
fn zero() -> Self::T { ... }
fn is_zero(x: &Self::T) -> bool
where Self::T: PartialEq { ... }
fn one() -> Self::T { ... }
fn is_one(x: &Self::T) -> bool
where Self::T: PartialEq { ... }
fn add(x: &Self::T, y: &Self::T) -> Self::T { ... }
fn mul(x: &Self::T, y: &Self::T) -> Self::T { ... }
fn try_matrix_product(
_a: &[Vec<Self::T>],
_b: &[Vec<Self::T>],
) -> Option<Vec<Vec<Self::T>>> { ... }
fn dot_product(x: &[Self::T], y: &[Self::T]) -> Self::T { ... }
fn add_scaled_assign(x: &mut [Self::T], y: &[Self::T], a: &Self::T) { ... }
fn add_assign(x: &mut Self::T, y: &Self::T) { ... }
fn mul_assign(x: &mut Self::T, y: &Self::T) { ... }
}Required Associated Types§
type T: Clone
type Additive: AbelianMonoid<T = Self::T>
type Multiplicative: Monoid<T = Self::T>
Provided Methods§
Sourcefn zero() -> Self::T
fn zero() -> Self::T
additive identity: $0$
Examples found in repository?
More examples
crates/competitive/src/algebra/ring.rs (line 20)
16 fn is_zero(x: &Self::T) -> bool
17 where
18 Self::T: PartialEq,
19 {
20 *x == Self::zero()
21 }
22 /// multiplicative identity: $1$
23 fn one() -> Self::T {
24 <Self::Multiplicative as Unital>::unit()
25 }
26 fn is_one(x: &Self::T) -> bool
27 where
28 Self::T: PartialEq,
29 {
30 *x == Self::one()
31 }
32 /// additive operaion: $+$
33 fn add(x: &Self::T, y: &Self::T) -> Self::T {
34 <Self::Additive as Magma>::operate(x, y)
35 }
36 /// multiplicative operaion: $+$
37 fn mul(x: &Self::T, y: &Self::T) -> Self::T {
38 <Self::Multiplicative as Magma>::operate(x, y)
39 }
40
41 fn try_matrix_product(_a: &[Vec<Self::T>], _b: &[Vec<Self::T>]) -> Option<Vec<Vec<Self::T>>> {
42 None
43 }
44
45 fn dot_product(x: &[Self::T], y: &[Self::T]) -> Self::T {
46 assert_eq!(x.len(), y.len());
47 x.iter().zip(y).fold(Self::zero(), |mut sum, (x, y)| {
48 Self::add_assign(&mut sum, &Self::mul(x, y));
49 sum
50 })
51 }crates/competitive/src/math/matrix.rs (line 90)
87 pub fn zeros(shape: (usize, usize)) -> Self {
88 Self {
89 shape,
90 data: vec![vec![R::zero(); shape.1]; shape.0],
91 _marker: PhantomData,
92 }
93 }
94
95 pub fn eye(shape: (usize, usize)) -> Self {
96 let mut data = vec![vec![R::zero(); shape.1]; shape.0];
97 for (i, d) in data.iter_mut().enumerate().take(shape.1) {
98 d[i] = R::one();
99 }
100 Self {
101 shape,
102 data,
103 _marker: PhantomData,
104 }
105 }
106
107 pub fn transpose(&self) -> Self {
108 Self::new_with((self.shape.1, self.shape.0), |i, j| self[j][i].clone())
109 }
110
111 pub fn map<S, F>(&self, mut f: F) -> Matrix<S>
112 where
113 S: SemiRing,
114 F: FnMut(&R::T) -> S::T,
115 {
116 Matrix::<S>::new_with(self.shape, |i, j| f(&self[i][j]))
117 }
118
119 pub fn add_row_with(&mut self, mut f: impl FnMut(usize, usize) -> R::T) {
120 self.data
121 .push((0..self.shape.1).map(|j| f(self.shape.0, j)).collect());
122 self.shape.0 += 1;
123 }
124
125 pub fn add_col_with(&mut self, mut f: impl FnMut(usize, usize) -> R::T) {
126 for i in 0..self.shape.0 {
127 self.data[i].push(f(i, self.shape.1));
128 }
129 self.shape.1 += 1;
130 }
131
132 pub fn pairwise_assign<F>(&mut self, other: &Self, mut f: F)
133 where
134 F: FnMut(&mut R::T, &R::T),
135 {
136 assert_eq!(self.shape, other.shape);
137 for i in 0..self.shape.0 {
138 for j in 0..self.shape.1 {
139 f(&mut self[i][j], &other[i][j]);
140 }
141 }
142 }
143}
144
145#[derive(Debug)]
146pub struct SystemOfLinearEquationsSolution<R>
147where
148 R: Field<Additive: Invertible, Multiplicative: Invertible>,
149{
150 pub particular: Vec<R::T>,
151 pub basis: Vec<Vec<R::T>>,
152}
153
154impl<R> Matrix<R>
155where
156 R: Field<T: PartialEq, Additive: Invertible, Multiplicative: Invertible>,
157{
158 fn eliminate<const DETERMINANT: bool>(&mut self) -> (usize, R::T) {
159 let (n, m) = self.shape;
160 let mut rank = 0;
161 let mut determinant = R::one();
162 let mut negative = false;
163 for first in (0..m).step_by(64) {
164 if rank == n {
165 break;
166 }
167 let end = (first + 64).min(m);
168 let start = rank;
169 let mut pivots = Vec::new();
170 let mut panel = vec![vec![R::zero(); end - first]; end - first];
171 for col in first..end {
172 if panel[col - first][..col - first]
173 .iter()
174 .any(|x| !R::is_zero(x))
175 {
176 for row in &mut self.data[rank..] {
177 let value =
178 R::dot_product(&row[first..col], &panel[col - first][..col - first]);
179 R::sub_assign(&mut row[col], &value);
180 }
181 }
182 let Some(pivot) = (rank..n).find(|&i| !R::is_zero(&self[i][col])) else {
183 continue;
184 };
185 if pivot != rank {
186 self.data.swap(rank, pivot);
187 negative = !negative;
188 }
189 if DETERMINANT {
190 R::mul_assign(&mut determinant, &self[rank][col]);
191 }
192 let inv = R::inv(&self[rank][col]);
193 let row = &mut self.data[rank];
194 for c in col + 1..end {
195 let value = R::dot_product(&row[first..col], &panel[c - first][..col - first]);
196 R::sub_assign(&mut row[c], &value);
197 panel[c - first][col - first] = row[c].clone();
198 }
199 for row in &mut self.data[rank + 1..] {
200 R::mul_assign(&mut row[col], &inv);
201 }
202 pivots.push(col);
203 rank += 1;
204 if rank == n {
205 break;
206 }
207 }
208 for i in start..if rank - start < 32 { n } else { rank } {
209 let (upper, lower) = self.data.split_at_mut(i);
210 let row = &mut lower[0];
211 for (j, &col) in pivots[..(i - start).min(pivots.len())].iter().enumerate() {
212 if R::is_zero(&row[col]) {
213 continue;
214 }
215 let factor = R::neg(&row[col]);
216 R::add_scaled_assign(&mut row[end..], &upper[start + j][end..], &factor);
217 }
218 }
219 if rank < n
220 && end < m
221 && rank - start >= 32
222 && self.data[rank..]
223 .iter()
224 .any(|row| pivots.iter().any(|&col| !R::is_zero(&row[col])))
225 {
226 let lower = Self::new_with((n - rank, rank - start), |i, j| {
227 R::neg(&self[rank + i][pivots[j]])
228 });
229 let upper = Self::new_with((rank - start, m - end), |i, j| {
230 self[start + i][end + j].clone()
231 });
232 let update = &lower * &upper;
233 for (row, update) in self.data[rank..].iter_mut().zip(update.data) {
234 for (x, y) in row[end..].iter_mut().zip(update) {
235 R::add_assign(x, &y);
236 }
237 }
238 }
239 for (i, &col) in pivots.iter().enumerate() {
240 for row in &mut self.data[start + i + 1..] {
241 row[col] = R::zero();
242 }
243 }
244 if DETERMINANT && rank < end {
245 return (rank, R::zero());
246 }
247 }
248 if DETERMINANT && negative {
249 determinant = R::neg(&determinant);
250 }
251 (rank, determinant)
252 }
253
254 /// f: (row, pivot_row, col)
255 pub fn row_reduction_with<F>(&mut self, normalize: bool, mut f: F)
256 where
257 F: FnMut(usize, usize, usize),
258 {
259 let (n, m) = self.shape;
260 let mut c = 0;
261 for r in 0..n {
262 loop {
263 if c >= m {
264 return;
265 }
266 if let Some(pivot) = (r..n).find(|&p| !R::is_zero(&self[p][c])) {
267 f(r, pivot, c);
268 self.data.swap(r, pivot);
269 break;
270 };
271 c += 1;
272 }
273 let d = R::inv(&self[r][c]);
274 if normalize {
275 for value in &mut self[r][c..m] {
276 R::mul_assign(value, &d);
277 }
278 }
279 for i in (0..n).filter(|&i| i != r) {
280 let mut e = self[i][c].clone();
281 if !normalize {
282 R::mul_assign(&mut e, &d);
283 }
284 for j in c..m {
285 let e = R::mul(&e, &self[r][j]);
286 R::sub_assign(&mut self[i][j], &e);
287 }
288 }
289 c += 1;
290 }
291 }
292
293 pub fn row_reduction(&mut self, normalize: bool) {
294 self.row_reduction_with(normalize, |_, _, _| {});
295 }
296
297 pub fn rank(&mut self) -> usize {
298 self.eliminate::<false>().0
299 }
300
301 pub fn determinant(&mut self) -> R::T {
302 assert_eq!(self.shape.0, self.shape.1);
303 self.eliminate::<true>().1
304 }
305
306 pub fn solve_system_of_linear_equations(
307 &self,
308 b: &[R::T],
309 ) -> Option<SystemOfLinearEquationsSolution<R>> {
310 assert_eq!(self.shape.0, b.len());
311 let m = self.shape.1;
312 let mut a = Self::new_with((self.shape.0, m + 1), |i, j| {
313 if j == m {
314 b[i].clone()
315 } else {
316 self[i][j].clone()
317 }
318 });
319 let rank = a.eliminate::<false>().0;
320 let mut pivots = Vec::with_capacity(rank);
321 let mut b = Vec::with_capacity(rank);
322 for row in &a.data[..rank] {
323 let c = row.iter().position(|x| !R::is_zero(x)).unwrap();
324 if c == m {
325 return None;
326 }
327 pivots.push(c);
328 b.push(row[m].clone());
329 }
330
331 let mut free = Vec::with_capacity(m - rank);
332 let mut pivot = 0;
333 for c in 0..m {
334 if pivot < rank && pivots[pivot] == c {
335 pivot += 1;
336 } else {
337 free.push(c);
338 }
339 }
340 let mut coefficients: Vec<Vec<_>> = (0..rank)
341 .map(|i| free.iter().map(|&c| a[i][c].clone()).collect())
342 .collect();
343 for k in (0..rank).rev() {
344 let c = pivots[k];
345 let inv = R::inv(&a[k][c]);
346 R::mul_assign(&mut b[k], &inv);
347 let pivot_b = b[k].clone();
348 let (upper, lower) = coefficients.split_at_mut(k);
349 let pivot_coefficients = &mut lower[0];
350 for x in pivot_coefficients.iter_mut() {
351 R::mul_assign(x, &inv);
352 }
353 for ((row, value), coefficients) in a.data[..k].iter_mut().zip(&mut b[..k]).zip(upper) {
354 if R::is_zero(&row[c]) {
355 continue;
356 }
357 let factor = row[c].clone();
358 row[c] = R::zero();
359 R::sub_assign(value, &R::mul(&factor, &pivot_b));
360 R::add_scaled_assign(coefficients, pivot_coefficients, &R::neg(&factor));
361 }
362 }
363
364 let mut particular = vec![R::zero(); m];
365 for i in 0..rank {
366 particular[pivots[i]] = b[i].clone();
367 }
368 let mut basis = Vec::with_capacity(free.len());
369 for (j, &c) in free.iter().enumerate() {
370 let mut vector = vec![R::zero(); m];
371 vector[c] = R::one();
372 for i in 0..rank {
373 vector[pivots[i]] = R::neg(&coefficients[i][j]);
374 }
375 basis.push(vector);
376 }
377 Some(SystemOfLinearEquationsSolution { particular, basis })
378 }
379
380 pub fn inverse(&self) -> Option<Matrix<R>> {
381 assert_eq!(self.shape.0, self.shape.1);
382 let n = self.shape.0;
383 if n >= 64 {
384 let m = n / 2;
385 let a = Self::new_with((m, m), |i, j| self[i][j].clone());
386 if let Some(mut ai) = a.inverse() {
387 let b = Self::new_with((m, n - m), |i, j| self[i][j + m].clone());
388 let c = Self::new_with((n - m, m), |i, j| self[i + m][j].clone());
389 let mut d = Self::new_with((n - m, n - m), |i, j| self[i + m][j + m].clone());
390 let u = &ai * &b;
391 let v = &c * &ai;
392 d -= &v * &b;
393 let di = d.inverse()?;
394 let r = &u * &di;
395 let t = &di * &v;
396 ai += &r * &v;
397 let mut inverse = Self::zeros((n, n));
398 for i in 0..m {
399 inverse[i][..m].clone_from_slice(&ai[i]);
400 for (x, y) in inverse[i][m..].iter_mut().zip(&r[i]) {
401 *x = R::neg(y);
402 }
403 }
404 for i in m..n {
405 for (x, y) in inverse[i][..m].iter_mut().zip(&t[i - m]) {
406 *x = R::neg(y);
407 }
408 inverse[i][m..].clone_from_slice(&di[i - m]);
409 }
410 return Some(inverse);
411 }
412 }
413 let mut a = self.clone();
414 let mut inverse = Self::eye((n, n));
415 let mut ranges: Vec<_> = (0..n).map(|i| (i, i + 1)).collect();
416 for r in 0..n {
417 let pivot = (r..n).find(|&i| !R::is_zero(&a[i][r]))?;
418 a.data.swap(r, pivot);
419 inverse.data.swap(r, pivot);
420 ranges.swap(r, pivot);
421
422 let d = R::inv(&a[r][r]);
423 for x in &mut a[r][r..] {
424 R::mul_assign(x, &d);
425 }
426 let (left, right) = ranges[r];
427 for x in &mut inverse[r][left..right] {
428 R::mul_assign(x, &d);
429 }
430
431 let (a_upper, a_lower) = a.data.split_at_mut(r + 1);
432 let pivot_a = &a_upper[r];
433 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r + 1);
434 let pivot_inverse = &inverse_upper[r];
435 let (ranges_upper, ranges_lower) = ranges.split_at_mut(r + 1);
436 let (left, right) = ranges_upper[r];
437 for ((a, inverse), range) in a_lower.iter_mut().zip(inverse_lower).zip(ranges_lower) {
438 if R::is_zero(&a[r]) {
439 continue;
440 }
441 let e = a[r].clone();
442 a[r] = R::zero();
443 R::add_scaled_assign(&mut a[(r + 1)..], &pivot_a[(r + 1)..], &R::neg(&e));
444 R::add_scaled_assign(
445 &mut inverse[left..right],
446 &pivot_inverse[left..right],
447 &R::neg(&e),
448 );
449 range.0 = range.0.min(left);
450 range.1 = range.1.max(right);
451 }
452 }
453 for r in (0..n).rev() {
454 let (left, right) = ranges[r];
455 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r);
456 let pivot_inverse = &inverse_lower[0];
457 let (ranges_upper, _) = ranges.split_at_mut(r);
458 for ((a, inverse), range) in a.data[..r].iter_mut().zip(inverse_upper).zip(ranges_upper)
459 {
460 if R::is_zero(&a[r]) {
461 continue;
462 }
463 let e = a[r].clone();
464 a[r] = R::zero();
465 R::add_scaled_assign(
466 &mut inverse[left..right],
467 &pivot_inverse[left..right],
468 &R::neg(&e),
469 );
470 range.0 = range.0.min(left);
471 range.1 = range.1.max(right);
472 }
473 }
474 Some(inverse)
475 }
476
477 pub fn characteristic_polynomial(&mut self) -> Vec<R::T> {
478 let n = self.shape.0;
479 if n == 0 {
480 return vec![R::one()];
481 }
482 assert!(self.data.iter().all(|a| a.len() == n));
483 for j in 0..(n - 1) {
484 if let Some(x) = ((j + 1)..n).find(|&x| !R::is_zero(&self[x][j])) {
485 self.data.swap(j + 1, x);
486 self.data.iter_mut().for_each(|a| a.swap(j + 1, x));
487 let inv = R::inv(&self[j + 1][j]);
488 let mut v = vec![];
489 let src = std::mem::take(&mut self[j + 1]);
490 for a in self.data[(j + 2)..].iter_mut() {
491 let mul = R::mul(&a[j], &inv);
492 R::add_scaled_assign(&mut a[j..], &src[j..], &R::neg(&mul));
493 v.push(mul);
494 }
495 self[j + 1] = src;
496 for a in self.data.iter_mut() {
497 let v = R::dot_product(&a[(j + 2)..], &v);
498 R::add_assign(&mut a[j + 1], &v);
499 }
500 }
501 }
502 // dp[k][j - k] stores [x^k] det(xI - A[..j, ..j]).
503 let mut dp: Vec<Vec<R::T>> = (0..=n).map(|i| Vec::with_capacity(n + 1 - i)).collect();
504 dp[0].push(R::one());
505 for i in 0..n {
506 let mut c = vec![R::zero(); i + 1];
507 c[i] = R::neg(&self[i][i]);
508 let mut mul = R::one();
509 for j in (0..i).rev() {
510 mul = R::mul(&mul, &self[j + 1][j]);
511 c[j] = R::neg(&R::mul(&mul, &self[j][i]));
512 }
513 for k in (0..=i).rev() {
514 let mut value = R::dot_product(&dp[k], &c[k..]);
515 if k > 0 {
516 R::add_assign(&mut value, dp[k - 1].last().unwrap());
517 }
518 dp[k].push(value);
519 }
520 dp[i + 1].push(R::one());
521 }
522 dp.into_iter().map(|mut c| c.pop().unwrap()).collect()
523 }
524}
525
526impl<R> Index<usize> for Matrix<R>
527where
528 R: SemiRing,
529{
530 type Output = Vec<R::T>;
531 fn index(&self, index: usize) -> &Self::Output {
532 &self.data[index]
533 }
534}
535
536impl<R> IndexMut<usize> for Matrix<R>
537where
538 R: SemiRing,
539{
540 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
541 &mut self.data[index]
542 }
543}
544
545impl<R> Index<(usize, usize)> for Matrix<R>
546where
547 R: SemiRing,
548{
549 type Output = R::T;
550 fn index(&self, index: (usize, usize)) -> &Self::Output {
551 &self.data[index.0][index.1]
552 }
553}
554
555impl<R> IndexMut<(usize, usize)> for Matrix<R>
556where
557 R: SemiRing,
558{
559 fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
560 &mut self.data[index.0][index.1]
561 }
562}
563
564macro_rules! impl_matrix_pairwise_binop {
565 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident $(where [$($clauses:tt)*])?) => {
566 impl<R> $imp_assign for Matrix<R>
567 where
568 R: SemiRing,
569 $($($clauses)*)?
570 {
571 fn $method_assign(&mut self, rhs: Self) {
572 self.pairwise_assign(&rhs, |a, b| R::$method_assign(a, b));
573 }
574 }
575 impl<R> $imp_assign<&Matrix<R>> for Matrix<R>
576 where
577 R: SemiRing,
578 $($($clauses)*)?
579 {
580 fn $method_assign(&mut self, rhs: &Self) {
581 self.pairwise_assign(rhs, |a, b| R::$method_assign(a, b));
582 }
583 }
584 impl<R> $imp for Matrix<R>
585 where
586 R: SemiRing,
587 $($($clauses)*)?
588 {
589 type Output = Matrix<R>;
590 fn $method(mut self, rhs: Self) -> Self::Output {
591 self.$method_assign(rhs);
592 self
593 }
594 }
595 impl<R> $imp<&Matrix<R>> for Matrix<R>
596 where
597 R: SemiRing,
598 $($($clauses)*)?
599 {
600 type Output = Matrix<R>;
601 fn $method(mut self, rhs: &Self) -> Self::Output {
602 self.$method_assign(rhs);
603 self
604 }
605 }
606 impl<R> $imp<Matrix<R>> for &Matrix<R>
607 where
608 R: SemiRing,
609 $($($clauses)*)?
610 {
611 type Output = Matrix<R>;
612 fn $method(self, mut rhs: Matrix<R>) -> Self::Output {
613 rhs.pairwise_assign(self, |a, b| *a = R::$method(b, a));
614 rhs
615 }
616 }
617 impl<R> $imp<&Matrix<R>> for &Matrix<R>
618 where
619 R: SemiRing,
620 $($($clauses)*)?
621 {
622 type Output = Matrix<R>;
623 fn $method(self, rhs: &Matrix<R>) -> Self::Output {
624 let mut this = self.clone();
625 this.$method_assign(rhs);
626 this
627 }
628 }
629 };
630}
631
632impl_matrix_pairwise_binop!(Add, add, AddAssign, add_assign);
633impl_matrix_pairwise_binop!(Sub, sub, SubAssign, sub_assign where [R: SemiRing<Additive: Invertible>]);
634
635impl<R> Mul for Matrix<R>
636where
637 R: SemiRing,
638{
639 type Output = Matrix<R>;
640 fn mul(self, rhs: Self) -> Self::Output {
641 (&self).mul(&rhs)
642 }
643}
644impl<R> Mul<&Matrix<R>> for Matrix<R>
645where
646 R: SemiRing,
647{
648 type Output = Matrix<R>;
649 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
650 (&self).mul(rhs)
651 }
652}
653impl<R> Mul<Matrix<R>> for &Matrix<R>
654where
655 R: SemiRing,
656{
657 type Output = Matrix<R>;
658 fn mul(self, rhs: Matrix<R>) -> Self::Output {
659 self.mul(&rhs)
660 }
661}
662impl<R> Mul<&Matrix<R>> for &Matrix<R>
663where
664 R: SemiRing,
665{
666 type Output = Matrix<R>;
667 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
668 assert_eq!(self.shape.1, rhs.shape.0);
669 if let Some(data) = R::try_matrix_product(&self.data, &rhs.data) {
670 return Matrix::from_vec(data);
671 }
672 let rhs = rhs.transpose();
673 Matrix::new_with((self.shape.0, rhs.shape.0), |i, j| {
674 R::dot_product(&self[i], &rhs[j])
675 })
676 }
677}
678
679fn strassen_rec<R: Ring>(
680 a: &[R::T],
681 b: &[R::T],
682 c: &mut [R::T],
683 shape: (usize, usize, usize),
684 stride_a: usize,
685 stride_b: usize,
686) {
687 let (n, m, p) = shape;
688 fn add_block<R: Ring>(
689 a: &[R::T],
690 b: &[R::T],
691 out: &mut [R::T],
692 n: usize,
693 stride_a: usize,
694 stride_b: usize,
695 ) {
696 for ((a, b), c) in a
697 .chunks(stride_a)
698 .zip(b.chunks(stride_b))
699 .zip(out.chunks_exact_mut(n))
700 {
701 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
702 *c = R::add(a, b);
703 }
704 }
705 }
706
707 fn sub_block<R: Ring>(
708 a: &[R::T],
709 b: &[R::T],
710 out: &mut [R::T],
711 n: usize,
712 stride_a: usize,
713 stride_b: usize,
714 ) {
715 for ((a, b), c) in a
716 .chunks(stride_a)
717 .zip(b.chunks(stride_b))
718 .zip(out.chunks_exact_mut(n))
719 {
720 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
721 *c = R::sub(a, b);
722 }
723 }
724 }
725
726 if n.min(m).min(p) <= 128 {
727 let transposed: Vec<_> = (0..p)
728 .flat_map(|j| (0..m).map(move |i| b[i * stride_b + j].clone()))
729 .collect();
730 for (a, c) in a.chunks(stride_a).zip(c.chunks_exact_mut(p)) {
731 for (b, c) in transposed.chunks_exact(m).zip(c) {
732 *c = R::dot_product(&a[..m], b);
733 }
734 }
735 return;
736 }
737 let (h, k, w) = (n / 2, m / 2, p / 2);
738 let a11 = 0;
739 let a12 = k;
740 let a21 = h * stride_a;
741 let a22 = a21 + k;
742 let b11 = 0;
743 let b12 = w;
744 let b21 = k * stride_b;
745 let b22 = b21 + w;
746
747 let block = h * w;
748 let mut buf = vec![R::zero(); h * k + k * w + block * 7];
749 let (s1, rest) = buf.split_at_mut(h * k);
750 let (s2, m_buf) = rest.split_at_mut(k * w);
751 let (m1, rest) = m_buf.split_at_mut(block);
752 let (m2, rest) = rest.split_at_mut(block);
753 let (m3, rest) = rest.split_at_mut(block);
754 let (m4, rest) = rest.split_at_mut(block);
755 let (m5, rest) = rest.split_at_mut(block);
756 let (m6, m7) = rest.split_at_mut(block);
757
758 // (A11 + A22)(B11 + B22)
759 add_block::<R>(&a[a11..], &a[a22..], s1, k, stride_a, stride_a);
760 add_block::<R>(&b[b11..], &b[b22..], s2, w, stride_b, stride_b);
761 strassen_rec::<R>(s1, s2, m1, (h, k, w), k, w);
762
763 // (A21 + A22) B11
764 add_block::<R>(&a[a21..], &a[a22..], s1, k, stride_a, stride_a);
765 strassen_rec::<R>(s1, &b[b11..], m2, (h, k, w), k, stride_b);
766
767 // A11 (B12 - B22)
768 sub_block::<R>(&b[b12..], &b[b22..], s2, w, stride_b, stride_b);
769 strassen_rec::<R>(&a[a11..], s2, m3, (h, k, w), stride_a, w);
770
771 // A22 (B21 - B11)
772 sub_block::<R>(&b[b21..], &b[b11..], s2, w, stride_b, stride_b);
773 strassen_rec::<R>(&a[a22..], s2, m4, (h, k, w), stride_a, w);
774
775 // (A11 + A12) B22
776 add_block::<R>(&a[a11..], &a[a12..], s1, k, stride_a, stride_a);
777 strassen_rec::<R>(s1, &b[b22..], m5, (h, k, w), k, stride_b);
778
779 // (A21 - A11)(B11 + B12)
780 sub_block::<R>(&a[a21..], &a[a11..], s1, k, stride_a, stride_a);
781 add_block::<R>(&b[b11..], &b[b12..], s2, w, stride_b, stride_b);
782 strassen_rec::<R>(s1, s2, m6, (h, k, w), k, w);
783
784 // (A12 - A22)(B21 + B22)
785 sub_block::<R>(&a[a12..], &a[a22..], s1, k, stride_a, stride_a);
786 add_block::<R>(&b[b21..], &b[b22..], s2, w, stride_b, stride_b);
787 strassen_rec::<R>(s1, s2, m7, (h, k, w), k, w);
788
789 let c11 = 0;
790 let c12 = w;
791 let c21 = h * p;
792 let c22 = c21 + w;
793 for ((((m1, m4), m5), m7), c) in m1
794 .iter()
795 .zip(m4.iter())
796 .zip(m5.iter())
797 .zip(m7.iter())
798 .zip(c[c11..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
799 {
800 *c = R::add(m1, m4);
801 R::sub_assign(c, m5);
802 R::add_assign(c, m7);
803 }
804 for ((m3, m5), c) in m3
805 .iter()
806 .zip(m5.iter())
807 .zip(c[c12..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
808 {
809 *c = R::add(m3, m5);
810 }
811 for ((m2, m4), c) in m2
812 .iter()
813 .zip(m4.iter())
814 .zip(c[c21..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
815 {
816 *c = R::add(m2, m4);
817 }
818 for ((((m1, m2), m3), m6), c) in m1
819 .iter()
820 .zip(m2.iter())
821 .zip(m3.iter())
822 .zip(m6.iter())
823 .zip(c[c22..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
824 {
825 *c = R::sub(m1, m2);
826 R::add_assign(c, m3);
827 R::add_assign(c, m6);
828 }
829}
830
831impl<R> Matrix<R>
832where
833 R: Ring,
834{
835 pub fn mul_strassen(&self, rhs: &Matrix<R>) -> Matrix<R> {
836 assert_eq!(self.shape.1, rhs.shape.0);
837 if let Some(data) = R::try_matrix_product(&self.data, &rhs.data) {
838 return Matrix::from_vec(data);
839 }
840 let (n, m) = self.shape;
841 let p = rhs.shape.1;
842 if n == 0 || m == 0 || p == 0 {
843 return Matrix::zeros((n, p));
844 }
845 let split = n.min(m).min(p).div_ceil(128).next_power_of_two();
846 if split <= 2 {
847 return self * rhs;
848 }
849 let rows = n.div_ceil(split) * split;
850 let inner = m.div_ceil(split) * split;
851 let cols = p.div_ceil(split) * split;
852 let mut a = vec![R::zero(); rows * inner];
853 for (a, data) in a.chunks_exact_mut(inner).zip(&self.data) {
854 a[..m].clone_from_slice(data);
855 }
856 let mut b = vec![R::zero(); inner * cols];
857 for (b, data) in b.chunks_exact_mut(cols).zip(&rhs.data) {
858 b[..p].clone_from_slice(data);
859 }
860 let mut c = vec![R::zero(); rows * cols];
861 strassen_rec::<R>(&a, &b, &mut c, (rows, inner, cols), inner, cols);
862 let mut res = Matrix::zeros((n, p));
863 for (data, c) in res.data.iter_mut().zip(c.chunks_exact(cols)) {
864 data.clone_from_slice(&c[..p]);
865 }
866 res
867 }crates/competitive/src/math/floor_sum.rs (line 174)
173 fn to_x() -> FloorSumData<R, X, Y> {
174 let mut dp = array![array![R::zero(); Y]; X];
175 dp[0][0] = R::one();
176 FloorSumData {
177 dp,
178 dx: R::one(),
179 dy: R::zero(),
180 _marker: PhantomData,
181 }
182 }
183 fn to_y() -> FloorSumData<R, X, Y> {
184 FloorSumData {
185 dp: array![array![R::zero(); Y]; X],
186 dx: R::zero(),
187 dy: R::one(),
188 _marker: PhantomData,
189 }
190 }
191}
192
193impl<R, const X: usize, const Y: usize> FloorSum<R, X, Y>
194where
195 R: Ring<Additive: Invertible>,
196{
197 fn offset(x: i64, y: i64) -> FloorSumData<R, X, Y> {
198 FloorSumData {
199 dp: array![array![R::zero(); Y]; X],
200 dx: R::Additive::signed_pow(R::one(), x),
201 dy: R::Additive::signed_pow(R::one(), y),
202 _marker: PhantomData,
203 }
204 }
205}
206
207impl<R, const X: usize, const Y: usize> Magma for FloorSum<R, X, Y>
208where
209 R: SemiRing,
210{
211 type T = FloorSumData<R, X, Y>;
212
213 fn operate(a: &Self::T, b: &Self::T) -> Self::T {
214 let mut a = a.clone();
215 let mut b = b.clone();
216 let mut pow_x = array![R::zero(); X];
217 let mut pow_y = array![R::zero(); Y];
218 pow_x[0] = R::one();
219 pow_y[0] = R::one();
220 for i in 1..X {
221 pow_x[i] = R::mul(&pow_x[i - 1], &a.dx);
222 }
223 for j in 1..Y {
224 pow_y[j] = R::mul(&pow_y[j - 1], &a.dy);
225 }
226 macro_rules! go {
227 ($N:ident) => {
228 let mut comb = array![array![R::zero(); $N]; $N];
229 comb[0][0] = R::one();
230 let mut i = 0;
231 while i + 1 < $N {
232 let mut j = 0;
233 while j <= i {
234 let x = comb[i][j].clone();
235 R::add_assign(&mut comb[i + 1][j], &x);
236 R::add_assign(&mut comb[i + 1][j + 1], &x);
237 j += 1;
238 }
239 i += 1;
240 }
241 for i in 0..X {
242 for j in (0..Y).rev() {
243 for k in j + 1..Y {
244 let mut x = b.dp[i][j].clone();
245 R::mul_assign(&mut x, &comb[k][j]);
246 R::mul_assign(&mut x, &pow_y[k - j]);
247 R::add_assign(&mut b.dp[i][k], &x);
248 }
249 }
250 }
251 for j in 0..Y {
252 for i in (0..X).rev() {
253 for k in i..X {
254 let mut x = b.dp[i][j].clone();
255 R::mul_assign(&mut x, &comb[k][i]);
256 R::mul_assign(&mut x, &pow_x[k - i]);
257 R::add_assign(&mut a.dp[k][j], &x);
258 }
259 }
260 }
261 };
262 }
263 if X <= Y {
264 go!(Y);
265 } else {
266 go!(X);
267 }
268 R::add_assign(&mut a.dx, &b.dx);
269 R::add_assign(&mut a.dy, &b.dy);
270 a
271 }
272}
273
274impl<R, const X: usize, const Y: usize> Unital for FloorSum<R, X, Y>
275where
276 R: SemiRing,
277{
278 fn unit() -> Self::T {
279 FloorSumData {
280 dp: array![array![R::zero(); Y]; X],
281 dx: R::zero(),
282 dy: R::zero(),
283 _marker: PhantomData,
284 }
285 }
286}
287
288impl<R, const X: usize, const Y: usize> Associative for FloorSum<R, X, Y> where R: SemiRing {}
289
290fn floor_monoid_product<M>(
291 mut x: M::T,
292 mut y: M::T,
293 mut n: u64,
294 mut a: u64,
295 mut b: u64,
296 mut m: u64,
297) -> M::T
298where
299 M: Monoid,
300{
301 let mut c = (a * n + b) / m;
302 let mut pre = M::unit();
303 let mut suf = M::unit();
304 loop {
305 let (p, q) = (a / m, b / m);
306 a %= m;
307 b %= m;
308 x = M::operate(&x, &M::pow(y.clone(), p));
309 pre = M::operate(&pre, &M::pow(y.clone(), q));
310 c -= p * n + q;
311 if c == 0 {
312 break;
313 }
314 let d = (m * c - b - 1) / a + 1;
315 suf = M::operate(&y, &M::operate(&M::pow(x.clone(), n - d), &suf));
316 b = m - b - 1 + a;
317 n = c - 1;
318 c = d;
319 swap(&mut m, &mut a);
320 swap(&mut x, &mut y);
321 }
322 x = M::pow(x.clone(), n);
323 M::operate(&M::operate(&pre, &x), &suf)
324}
325
326/// $$\sum_{i=0}^{n-1}i^X\left\lfloor\frac{a\times i+b}{m}\right\rfloor^Y$$
327pub fn floor_sum_polynomial<T, const X: usize, const Y: usize>(
328 n: u64,
329 a: u64,
330 b: u64,
331 m: u64,
332) -> [[T; Y]; X]
333where
334 T: DotProduct + One,
335{
336 debug_assert!(a == 0 || n < (u64::MAX - b) / a);
337 floor_monoid_product::<FloorSum<AddMulOperation<T>, X, Y>>(
338 FloorSum::<AddMulOperation<T>, X, Y>::to_x(),
339 FloorSum::<AddMulOperation<T>, X, Y>::to_y(),
340 n,
341 a,
342 b,
343 m,
344 )
345 .dp
346}
347
348/// $$\sum_{i=l}^{r-1}i^X\left\lfloor\frac{a\times i+b}{m}\right\rfloor^Y$$
349pub fn floor_sum_polynomial_i64<T, const X: usize, const Y: usize>(
350 l: i64,
351 r: i64,
352 a: i64,
353 b: i64,
354 m: u64,
355) -> [[T; Y]; X]
356where
357 T: DotProduct + One,
358 AddMulOperation<T>: SemiRing<T = T, Additive: Invertible>,
359{
360 assert!(l <= r);
361 assert!(m > 0);
362
363 if a < 0 {
364 let mut ans = floor_sum_polynomial_i64::<T, X, Y>(-r + 1, -l + 1, -a, b, m);
365 for ans in ans.iter_mut().skip(1).step_by(2) {
366 for ans in ans.iter_mut() {
367 *ans = AddMulOperation::<T>::neg(ans);
368 }
369 }
370 return ans;
371 }
372
373 let add_x = l;
374 let n = (r - l) as u64;
375 let b = a * add_x + b;
376
377 let add_y = b.div_euclid(m as i64);
378 let b = b.rem_euclid(m as i64);
379 assert!(a >= 0);
380 assert!(b >= 0);
381 let data = floor_monoid_product::<FloorSum<AddMulOperation<T>, X, Y>>(
382 FloorSum::<AddMulOperation<T>, X, Y>::to_x(),
383 FloorSum::<AddMulOperation<T>, X, Y>::to_y(),
384 n,
385 a as u64,
386 b as u64,
387 m,
388 );
389
390 let offset = FloorSum::<AddMulOperation<T>, X, Y>::offset(add_x, add_y);
391 FloorSum::<AddMulOperation<T>, X, Y>::operate(&offset, &data).dp
392}
393
394#[derive(Debug)]
395struct FloorPowerSum<R>
396where
397 R: SemiRing,
398{
399 x: R::T,
400 sum: R::T,
401}
402
403impl<R> Clone for FloorPowerSum<R>
404where
405 R: SemiRing,
406{
407 fn clone(&self) -> Self {
408 Self {
409 x: self.x.clone(),
410 sum: self.sum.clone(),
411 }
412 }
413}
414
415impl<R> FloorPowerSum<R>
416where
417 R: SemiRing,
418{
419 fn to_x(x: R::T) -> Self {
420 Self { x, sum: R::one() }
421 }
422 fn to_y(y: R::T) -> Self {
423 Self {
424 x: y,
425 sum: R::zero(),
426 }
427 }
428}
429
430impl<R> Magma for FloorPowerSum<R>
431where
432 R: SemiRing,
433{
434 type T = Self;
435
436 fn operate(a: &Self::T, b: &Self::T) -> Self::T {
437 Self {
438 x: R::mul(&a.x, &b.x),
439 sum: R::add(&a.sum, &R::mul(&a.x, &b.sum)),
440 }
441 }
442}
443
444impl<R> Unital for FloorPowerSum<R>
445where
446 R: SemiRing,
447{
448 fn unit() -> Self::T {
449 Self {
450 x: R::one(),
451 sum: R::zero(),
452 }
453 }crates/competitive/src/math/subset_convolve.rs (line 16)
12 fn ranked(t: Vec<R::T>, len: usize) -> (Vec<R::T>, usize) {
13 let width = len.trailing_zeros() as usize + 1;
14 let mut ranked = Vec::with_capacity(len * width);
15 advise_huge_pages(&mut ranked);
16 ranked.resize(len * width, R::zero());
17 for (i, value) in t.into_iter().enumerate() {
18 ranked[i * width + i.count_ones() as usize] = value;
19 }
20 (ranked, width)
21 }
22
23 fn diagonal(ranked: Vec<R::T>, width: usize) -> Vec<R::T> {
24 ranked
25 .chunks_exact(width)
26 .enumerate()
27 .map(|(i, row)| row[i.count_ones() as usize].clone())
28 .collect()
29 }
30
31 #[inline]
32 fn multiply_row(
33 x: &[R::T],
34 y: &[R::T],
35 right: &mut [R::T],
36 output: &mut [R::T],
37 rank: usize,
38 ) -> usize {
39 for (right, y) in right[..=rank].iter_mut().zip(y[..=rank].iter().rev()) {
40 right.clone_from(y);
41 }
42 let end = (rank * 2).min(x.len() - 1);
43 for (degree, output) in output.iter_mut().enumerate().take(end + 1).skip(rank) {
44 let first = degree - rank;
45 *output = R::dot_product(&x[first..=rank], &right[..=rank - first]);
46 }
47 end
48 }
49}
50
51impl<R> ConvolveSteps for SubsetConvolve<R>
52where
53 R: Ring<T: PartialEq, Additive: Invertible>,
54{
55 type T = Vec<R::T>;
56 type F = (Vec<R::T>, usize);
57
58 fn length(t: &Self::T) -> usize {
59 t.len()
60 }
61
62 fn transform(t: Self::T, len: usize) -> Self::F {
63 let (mut f, width) = Self::ranked(t, len);
64 let k = width - 1;
65 for bit in 0..k {
66 let half = 1 << bit;
67 for base in (0..len).step_by(half * 2) {
68 for lower in base..base + half {
69 let upper = lower + half;
70 let ranks = lower.count_ones() as usize + 1;
71 let (lower_rows, upper_rows) = f.split_at_mut(upper * width);
72 let lower_row = &lower_rows[lower * width..lower * width + ranks];
73 let upper_row = &mut upper_rows[..ranks];
74 for (upper, lower) in upper_row.iter_mut().zip(lower_row) {
75 R::add_assign(upper, lower);
76 }
77 }
78 }
79 }
80 (f, width)
81 }
82
83 fn inverse_transform((mut f, width): Self::F, len: usize) -> Self::T {
84 let k = width - 1;
85 for bit in 0..k {
86 let half = 1 << bit;
87 for base in (0..len).step_by(half * 2) {
88 for lower in base..base + half {
89 let upper = lower + half;
90 let rank = lower.count_ones() as usize;
91 let (lower_rows, upper_rows) = f.split_at_mut(upper * width);
92 let lower_row = &lower_rows[lower * width + rank..lower * width + width];
93 let upper_row = &mut upper_rows[rank..width];
94 for (upper, lower) in upper_row.iter_mut().zip(lower_row) {
95 R::sub_assign(upper, lower);
96 }
97 }
98 }
99 }
100 Self::diagonal(f, width)
101 }
102
103 fn multiply(f: &mut Self::F, g: &Self::F) {
104 let (f, width) = f;
105 let (g, _) = g;
106 let mut right = vec![R::zero(); *width];
107 let mut output = vec![R::zero(); *width];
108 for (i, f) in f.chunks_exact_mut(*width).enumerate() {
109 let rank = i.count_ones() as usize;
110 let g = &g[i * *width..(i + 1) * *width];
111 let end = Self::multiply_row(f, g, &mut right, &mut output, rank);
112 f[rank..=end].clone_from_slice(&output[rank..=end]);
113 }
114 }
115
116 fn convolve(a: Self::T, b: Self::T) -> Self::T {
117 assert_eq!(a.len(), b.len());
118 let len = a.len();
119 let same = a == b;
120 let (mut x, width) = Self::ranked(a, len);
121 let (mut y, _) = if same {
122 (x.clone(), width)
123 } else {
124 Self::ranked(b, len)
125 };
126 let mut right = vec![R::zero(); width];
127 let mut output = vec![R::zero(); width];
128 for i in 0..len {
129 for bit in (0..(i | len).trailing_zeros() as usize).rev() {
130 let half = width << bit;
131 let start = i * width;
132 let (lower, upper) = x[start..start + half * 2].split_at_mut(half);
133 for (upper, lower) in upper.iter_mut().zip(lower) {
134 R::add_assign(upper, lower);
135 }
136 let (lower, upper) = y[start..start + half * 2].split_at_mut(half);
137 for (upper, lower) in upper.iter_mut().zip(lower) {
138 R::add_assign(upper, lower);
139 }
140 }
141
142 let rank = i.count_ones() as usize;
143 let start = i * width;
144 let x_row = &x[start..start + width];
145 let y_row = &y[start..start + width];
146 output.fill(R::zero());
147 Self::multiply_row(x_row, y_row, &mut right, &mut output, rank);
148 x[start..start + width].clone_from_slice(&output);
149
150 for bit in 0..i.trailing_ones() as usize {
151 let end = (i + 1) * width;
152 let half = width << bit;
153 let (lower, upper) = x[end - half * 2..end].split_at_mut(half);
154 for (upper, lower) in upper.iter_mut().zip(lower) {
155 R::sub_assign(upper, lower);
156 }
157 }
158 }
159 Self::diagonal(x, width)
160 }crates/competitive/src/algorithm/automata_learning.rs (line 222)
210 fn behavior<I>(&self, input: I) -> Self::Output
211 where
212 I: IntoIterator<Item = usize>,
213 {
214 let mut weights = self.initial_weights.clone();
215 for x in input {
216 weights = &weights * &self.transitions[x];
217 }
218 let result = &weights * &self.final_weights;
219 if result.shape != (0, 0) {
220 result[0][0].clone()
221 } else {
222 F::zero()
223 }
224 }
225}
226
227impl<F> SerdeByteStr for WeightedFiniteAutomaton<F>
228where
229 F: Field<T: SerdeByteStr, Additive: Invertible, Multiplicative: Invertible>,
230{
231 fn serialize(&self, buf: &mut Vec<u8>) {
232 self.initial_weights.serialize(buf);
233 self.transitions.serialize(buf);
234 self.final_weights.serialize(buf);
235 }
236
237 fn deserialize<I>(iter: &mut I) -> Self
238 where
239 I: Iterator<Item = u8>,
240 {
241 let initial_weights = Matrix::deserialize(iter);
242 let transitions = Vec::deserialize(iter);
243 let final_weights = Matrix::deserialize(iter);
244 Self {
245 initial_weights,
246 transitions,
247 final_weights,
248 }
249 }
250}
251
252pub fn dense_sampling(sigma: usize, max_len: usize) -> impl Iterator<Item = Vec<usize>> {
253 assert_ne!(sigma, 0, "Sigma must be greater than 0");
254 let mut current = vec![];
255 once_with(Vec::new).chain(from_fn(move || {
256 let mut carry = true;
257 for i in (0..current.len()).rev() {
258 current[i] += 1;
259 if current[i] == sigma {
260 current[i] = 0;
261 } else {
262 carry = false;
263 break;
264 }
265 }
266 if carry {
267 current.push(0);
268 }
269 if current.len() <= max_len {
270 Some(current.to_vec())
271 } else {
272 None
273 }
274 }))
275}
276
277pub fn random_sampling(
278 sigma: usize,
279 len_spec: impl RandomSpec<usize>,
280 seconds: f64,
281) -> impl Iterator<Item = Vec<usize>> {
282 assert_ne!(sigma, 0, "Sigma must be greater than 0");
283 let now = Instant::now();
284 let mut rng = Xorshift::new();
285 from_fn(move || {
286 if now.elapsed().as_secs_f64() > seconds {
287 None
288 } else {
289 let n = rng.random(&len_spec);
290 Some(rng.random_iter(0..sigma).take(n).collect())
291 }
292 })
293}
294
295#[derive(Debug, Clone)]
296pub struct DfaLearning<A>
297where
298 A: BlackBoxAutomaton<Output = bool>,
299{
300 automaton: A,
301 prefixes: Vec<Vec<usize>>,
302 suffixes: Vec<Vec<usize>>,
303 table: Vec<BitSet>,
304 row_map: HashMap<BitSet, usize>,
305}
306
307impl<A> DfaLearning<A>
308where
309 A: BlackBoxAutomaton<Output = bool>,
310{
311 pub fn new(automaton: A) -> Self {
312 let mut this = Self {
313 automaton,
314 prefixes: vec![],
315 suffixes: vec![],
316 table: vec![],
317 row_map: HashMap::new(),
318 };
319 this.add_suffix(vec![]);
320 this.add_prefix(vec![]);
321 this
322 }
323 fn add_prefix(&mut self, prefix: Vec<usize>) -> usize {
324 let row: BitSet = self
325 .suffixes
326 .iter()
327 .map(|s| {
328 self.automaton
329 .behavior(prefix.iter().cloned().chain(s.iter().cloned()))
330 })
331 .collect();
332 *self.row_map.entry(row.clone()).or_insert_with(|| {
333 let idx = self.table.len();
334 self.table.push(row);
335 self.prefixes.push(prefix);
336 idx
337 })
338 }
339 fn add_suffix(&mut self, suffix: Vec<usize>) {
340 if self.suffixes.contains(&suffix) {
341 return;
342 }
343 for (prefix, table) in self.prefixes.iter_mut().zip(&mut self.table) {
344 table.push(
345 self.automaton
346 .behavior(prefix.iter().cloned().chain(suffix.iter().cloned())),
347 );
348 }
349 self.suffixes.push(suffix);
350 self.row_map.clear();
351 for (i_prefix, row) in self.table.iter().enumerate() {
352 self.row_map.insert(row.clone(), i_prefix);
353 }
354 }
355 pub fn construct_dfa(&mut self) -> DeterministicFiniteAutomaton {
356 let sigma = self.automaton.sigma();
357 let mut dfa = DeterministicFiniteAutomaton {
358 states: vec![],
359 initial_state: 0,
360 };
361 let mut i_prefix = 0;
362 while i_prefix < self.prefixes.len() {
363 let mut delta = vec![];
364 for x in 0..sigma {
365 let prefix: Vec<usize> =
366 self.prefixes[i_prefix].iter().cloned().chain([x]).collect();
367 let index = self.add_prefix(prefix);
368 delta.push(index);
369 }
370 dfa.states.push(DfaState {
371 delta,
372 accept: self.table[i_prefix].get(0),
373 });
374 i_prefix += 1;
375 }
376 dfa
377 }
378 pub fn train_sample(&mut self, dfa: &DeterministicFiniteAutomaton, sample: &[usize]) -> bool {
379 let expected = self.automaton.behavior(sample.iter().cloned());
380 if expected == dfa.behavior(sample.iter().cloned()) {
381 return false;
382 }
383 let n = sample.len();
384 let mut states: Vec<(usize, usize)> = Vec::with_capacity(n + 1);
385 let mut s = 0usize;
386 states.push((s, 0));
387 for (k, &x) in sample.iter().enumerate() {
388 s = dfa.states[s].delta[x];
389 states.push((s, k + 1));
390 }
391 let split = states.partition_point(|&(state, k)| {
392 self.automaton.behavior(
393 self.prefixes[state]
394 .iter()
395 .cloned()
396 .chain(sample[k..].iter().cloned()),
397 ) == expected
398 });
399 let new_prefix = sample[..split].to_vec();
400 let new_suffix = sample[split..].to_vec();
401 self.add_suffix(new_suffix);
402 self.add_prefix(new_prefix);
403 true
404 }
405 pub fn train(
406 &mut self,
407 samples: impl IntoIterator<Item = Vec<usize>>,
408 ) -> DeterministicFiniteAutomaton {
409 let mut dfa = self.construct_dfa();
410 for sample in samples {
411 if self.train_sample(&dfa, &sample) {
412 dfa = self.construct_dfa();
413 }
414 }
415 dfa
416 }
417}
418
419pub struct WfaLearning<F, A>
420where
421 F: Field<Additive: Invertible, Multiplicative: Invertible>,
422 A: BlackBoxAutomaton<Output = F::T>,
423{
424 automaton: A,
425 prefixes: Vec<Vec<usize>>,
426 suffixes: Vec<Vec<usize>>,
427 inv_h: Matrix<F>,
428 nh: Vec<Matrix<F>>,
429 wfa: WeightedFiniteAutomaton<F>,
430 _marker: PhantomData<fn() -> F>,
431}
432
433impl<F, A> Debug for WfaLearning<F, A>
434where
435 F: Field<T: Debug, Additive: Invertible, Multiplicative: Invertible>,
436 A: BlackBoxAutomaton<Output = F::T> + Debug,
437{
438 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
439 f.debug_struct("WfaLearning")
440 .field("automaton", &self.automaton)
441 .field("prefixes", &self.prefixes)
442 .field("suffixes", &self.suffixes)
443 .field("inv_h", &self.inv_h)
444 .field("nh", &self.nh)
445 .field("wfa", &self.wfa)
446 .finish()
447 }
448}
449
450impl<F, A> Clone for WfaLearning<F, A>
451where
452 F: Field<Additive: Invertible, Multiplicative: Invertible>,
453 A: BlackBoxAutomaton<Output = F::T> + Clone,
454{
455 fn clone(&self) -> Self {
456 Self {
457 automaton: self.automaton.clone(),
458 prefixes: self.prefixes.clone(),
459 suffixes: self.suffixes.clone(),
460 inv_h: self.inv_h.clone(),
461 nh: self.nh.clone(),
462 wfa: self.wfa.clone(),
463 _marker: self._marker,
464 }
465 }
466}
467
468impl<F, A> WfaLearning<F, A>
469where
470 F: Field<T: PartialEq, Additive: Invertible, Multiplicative: Invertible>,
471 A: BlackBoxAutomaton<Output = F::T>,
472{
473 pub fn new(automaton: A) -> Self {
474 let sigma = automaton.sigma();
475 Self {
476 automaton,
477 prefixes: vec![],
478 suffixes: vec![],
479 inv_h: Matrix::zeros((0, 0)),
480 nh: vec![Matrix::zeros((0, 0)); sigma],
481 wfa: WeightedFiniteAutomaton {
482 initial_weights: Matrix::zeros((1, 0)),
483 transitions: vec![Matrix::zeros((0, 0)); sigma],
484 final_weights: Matrix::zeros((0, 1)),
485 },
486 _marker: PhantomData,
487 }
488 }
489 pub fn wfa(&self) -> &WeightedFiniteAutomaton<F> {
490 &self.wfa
491 }
492 fn split_sample(&mut self, sample: &[usize]) -> Option<(Vec<usize>, Vec<usize>)> {
493 if self.prefixes.is_empty() && !F::is_zero(&self.automaton.behavior(sample.iter().cloned()))
494 {
495 return Some((vec![], sample.to_vec()));
496 }
497 let expected = self.automaton.behavior(sample.iter().cloned());
498 if expected == self.wfa.behavior(sample.iter().cloned()) {
499 return None;
500 }
501 let n = sample.len();
502 let dim = self.wfa.final_weights.shape.0;
503 let mut states: Vec<(Matrix<F>, usize)> = Vec::with_capacity(n + 1);
504 let mut v = self.wfa.final_weights.clone();
505 states.push((v.clone(), n));
506 for k in (0..n).rev() {
507 v = &self.wfa.transitions[sample[k]] * &v;
508 states.push((v.clone(), k));
509 }
510 states.reverse();
511 let split = states.partition_point(|(state, k)| {
512 (0..dim).any(|j| {
513 self.automaton.behavior(
514 self.prefixes[j]
515 .iter()
516 .cloned()
517 .chain(sample[*k..].iter().cloned()),
518 ) != state[j][0]
519 })
520 });
521 Some((sample[..split].to_vec(), sample[split..].to_vec()))
522 }
523 pub fn train_sample(&mut self, sample: &[usize]) -> bool {
524 let Some((prefix, suffix)) = self.split_sample(sample) else {
525 return false;
526 };
527 self.prefixes.push(prefix);
528 self.suffixes.push(suffix);
529 let n = self.inv_h.shape.0;
530 let prefix = &self.prefixes[n];
531 let suffix = &self.suffixes[n];
532 let u = Matrix::<F>::new_with((n, 1), |i, _| {
533 self.automaton.behavior(
534 self.prefixes[i]
535 .iter()
536 .cloned()
537 .chain(suffix.iter().cloned()),
538 )
539 });
540 let v = Matrix::<F>::new_with((1, n), |_, j| {
541 self.automaton.behavior(
542 prefix
543 .iter()
544 .cloned()
545 .chain(self.suffixes[j].iter().cloned()),
546 )
547 });
548 let w = Matrix::<F>::new_with((1, 1), |_, _| {
549 self.automaton
550 .behavior(prefix.iter().cloned().chain(suffix.iter().cloned()))
551 });
552 let t = &self.inv_h * &u;
553 let s = &v * &self.inv_h;
554 let d = F::inv(&(&w - &(&v * &t))[0][0]);
555 let dh = &t * &s;
556 for i in 0..n {
557 for j in 0..n {
558 F::add_assign(&mut self.inv_h[i][j], &F::mul(&dh[i][j], &d));
559 }
560 }
561 self.inv_h
562 .add_col_with(|i, _| F::neg(&F::mul(&t[i][0], &d)));
563 self.inv_h.add_row_with(|_, j| {
564 if j != n {
565 F::neg(&F::mul(&s[0][j], &d))
566 } else {
567 d.clone()
568 }
569 });
570
571 for (x, transition) in self.wfa.transitions.iter_mut().enumerate() {
572 let b = &(&self.nh[x] * &t) * &s;
573 for i in 0..n {
574 for j in 0..n {
575 F::add_assign(&mut transition[i][j], &F::mul(&b[i][j], &d));
576 }
577 }
578 }
579 for (x, nh) in self.nh.iter_mut().enumerate() {
580 nh.add_col_with(|i, j| {
581 self.automaton.behavior(
582 self.prefixes[i]
583 .iter()
584 .cloned()
585 .chain([x])
586 .chain(self.suffixes[j].iter().cloned()),
587 )
588 });
589 nh.add_row_with(|i, j| {
590 self.automaton.behavior(
591 self.prefixes[i]
592 .iter()
593 .cloned()
594 .chain([x])
595 .chain(self.suffixes[j].iter().cloned()),
596 )
597 });
598 }
599 self.wfa
600 .initial_weights
601 .add_col_with(|_, _| if n == 0 { F::one() } else { F::zero() });
602 self.wfa
603 .final_weights
604 .add_row_with(|_, _| self.automaton.behavior(prefix.iter().cloned()));
605 for (x, transition) in self.wfa.transitions.iter_mut().enumerate() {
606 transition.add_col_with(|_, _| F::zero());
607 transition.add_row_with(|_, _| F::zero());
608 for i in 0..=n {
609 for j in 0..=n {
610 if i == n || j == n {
611 for k in 0..=n {
612 if i != n && j != n && k != n {
613 continue;
614 }
615 F::add_assign(
616 &mut transition[i][k],
617 &F::mul(&self.nh[x][i][j], &self.inv_h[j][k]),
618 );
619 }
620 } else {
621 let k = n;
622 F::add_assign(
623 &mut transition[i][k],
624 &F::mul(&self.nh[x][i][j], &self.inv_h[j][k]),
625 );
626 }
627 }
628 }
629 }
630 true
631 }
632 pub fn train(&mut self, samples: impl IntoIterator<Item = Vec<usize>>) {
633 for sample in samples {
634 self.train_sample(&sample);
635 }
636 }
637 pub fn batch_train(&mut self, samples: impl IntoIterator<Item = Vec<usize>>) {
638 let mut prefix_set: HashSet<_> = self.prefixes.iter().cloned().collect();
639 let mut suffix_set: HashSet<_> = self.suffixes.iter().cloned().collect();
640 for sample in samples {
641 if prefix_set.insert(sample.to_vec()) {
642 self.prefixes.push(sample.to_vec());
643 }
644 if suffix_set.insert(sample.to_vec()) {
645 self.suffixes.push(sample);
646 }
647 }
648 let mut h = Matrix::<F>::new_with((self.prefixes.len(), self.suffixes.len()), |i, j| {
649 self.automaton.behavior(
650 self.prefixes[i]
651 .iter()
652 .cloned()
653 .chain(self.suffixes[j].iter().cloned()),
654 )
655 });
656 if !self.prefixes.is_empty() && !self.suffixes.is_empty() && F::is_zero(&h[0][0]) {
657 for j in 1..self.suffixes.len() {
658 if !F::is_zero(&h[0][j]) {
659 self.suffixes.swap(0, j);
660 for row in &mut h.data {
661 row.swap(0, j);
662 }
663 break;
664 }
665 }
666 }
667 let mut row_id: Vec<usize> = (0..h.shape.0).collect();
668 let mut pivots = vec![];
669 h.row_reduction_with(false, |r, p, c| {
670 row_id.swap(r, p);
671 pivots.push((row_id[r], c));
672 });
673 let mut new_prefixes = vec![];
674 let mut new_suffixes = vec![];
675 for (i, j) in pivots {
676 new_prefixes.push(self.prefixes[i].clone());
677 new_suffixes.push(self.suffixes[j].clone());
678 }
679 self.prefixes = new_prefixes;
680 self.suffixes = new_suffixes;
681 assert_eq!(self.prefixes.len(), self.suffixes.len());
682 let n = self.prefixes.len();
683 let h = Matrix::<F>::new_with((n, n), |i, j| {
684 self.automaton.behavior(
685 self.prefixes[i]
686 .iter()
687 .cloned()
688 .chain(self.suffixes[j].iter().cloned()),
689 )
690 });
691 self.inv_h = h.inverse().expect("Hankel matrix must be invertible");
692 self.wfa = WeightedFiniteAutomaton::<F> {
693 initial_weights: Matrix::new_with((1, n), |_, j| {
694 if self.prefixes[j].is_empty() {
695 F::one()
696 } else {
697 F::zero()
698 }
699 }),
700 transitions: (0..self.automaton.sigma())
701 .map(|x| {
702 &Matrix::new_with((n, n), |i, j| {
703 self.automaton.behavior(
704 self.prefixes[i]
705 .iter()
706 .cloned()
707 .chain([x])
708 .chain(self.suffixes[j].iter().cloned()),
709 )
710 }) * &self.inv_h
711 })
712 .collect(),
713 final_weights: Matrix::new_with((n, 1), |i, _| {
714 self.automaton.behavior(self.prefixes[i].iter().cloned())
715 }),
716 };
717 }Additional examples can be found in:
Sourcefn is_zero(x: &Self::T) -> bool
fn is_zero(x: &Self::T) -> bool
Examples found in repository?
crates/competitive/src/math/black_box_matrix.rs (line 81)
72 pub fn new_with<F>(shape: (usize, usize), f: F) -> Self
73 where
74 R: SemiRing<T: PartialEq>,
75 F: Fn(usize, usize) -> R::T,
76 {
77 let mut nonzero = vec![];
78 for i in 0..shape.0 {
79 for j in 0..shape.1 {
80 let v = f(i, j);
81 if !R::is_zero(&v) {
82 nonzero.push((i, j, v));
83 }
84 }
85 }
86 Self { shape, nonzero }
87 }
88 pub fn from_nonzero(shape: (usize, usize), nonzero: Vec<(usize, usize, R::T)>) -> Self {
89 Self { shape, nonzero }
90 }
91}
92
93impl<R> SparseMatrix<R>
94where
95 R: Field<T: PartialEq, Additive: Invertible, Multiplicative: Invertible>,
96{
97 pub fn determinant(&self) -> R::T {
98 assert_eq!(self.shape.0, self.shape.1);
99 let n = self.shape.0;
100 let mut columns = vec![Vec::<(usize, R::T)>::new(); n];
101 for &(i, j, ref value) in &self.nonzero {
102 columns[j].push((i, value.clone()));
103 }
104 let mut degrees = vec![0; n];
105 for column in &mut columns {
106 column.sort_unstable_by_key(|&(i, _)| i);
107 let mut merged: Vec<(usize, R::T)> = Vec::with_capacity(column.len());
108 for (i, value) in column.drain(..) {
109 if let Some((_, x)) = merged.last_mut().filter(|(last, _)| *last == i) {
110 R::add_assign(x, &value);
111 } else {
112 merged.push((i, value));
113 }
114 }
115 merged.retain(|(i, value)| {
116 if R::is_zero(value) {
117 false
118 } else {
119 degrees[*i] += 1;
120 true
121 }
122 });
123 *column = merged;
124 }
125 let mut order: Vec<_> = (0..n).collect();
126 order.sort_unstable_by_key(|&j| columns[j].len());
127 let mut lower: Vec<Vec<(usize, R::T)>> = Vec::with_capacity(n);
128 let mut pivots: Vec<Option<usize>> = vec![None; n];
129 let mut x = vec![R::zero(); n];
130 let mut seen = vec![0; n];
131 let mut stack = Vec::new();
132 let mut support = Vec::new();
133 let mut determinant = R::one();
134 for (k, &j) in order.iter().enumerate() {
135 support.clear();
136 for &(i, _) in &columns[j] {
137 if seen[i] == k + 1 {
138 continue;
139 }
140 seen[i] = k + 1;
141 x[i] = R::zero();
142 stack.push((i, 0));
143 while let Some((i, next)) = stack.last_mut() {
144 if let Some(pivot) = pivots[*i]
145 && *next < lower[pivot].len()
146 {
147 let row = lower[pivot][*next].0;
148 *next += 1;
149 if seen[row] != k + 1 {
150 seen[row] = k + 1;
151 x[row] = R::zero();
152 stack.push((row, 0));
153 }
154 continue;
155 }
156 support.push(*i);
157 stack.pop();
158 }
159 }
160 for &(i, ref value) in &columns[j] {
161 x[i] = value.clone();
162 }
163 let mut pivot = None;
164 for &i in support.iter().rev() {
165 if let Some(p) = pivots[i] {
166 let factor = R::neg(&x[i]);
167 for &(row, ref value) in &lower[p] {
168 R::add_assign(&mut x[row], &R::mul(&factor, value));
169 }
170 } else if !R::is_zero(&x[i]) && pivot.is_none_or(|p| degrees[i] < degrees[p]) {
171 pivot = Some(i);
172 }
173 }
174 let Some(pivot) = pivot else { return R::zero() };
175 R::mul_assign(&mut determinant, &x[pivot]);
176 let inv = R::inv(&x[pivot]);
177 pivots[pivot] = Some(k);
178 lower.push(
179 support
180 .iter()
181 .filter(|&&i| pivots[i].is_none() && !R::is_zero(&x[i]))
182 .map(|&i| (i, R::mul(&x[i], &inv)))
183 .collect(),
184 );
185 }
186 for mut permutation in [order, pivots.into_iter().map(Option::unwrap).collect()] {
187 for i in 0..n {
188 while permutation[i] != i {
189 let j = permutation[i];
190 permutation.swap(i, j);
191 determinant = R::neg(&determinant);
192 }
193 }
194 }
195 determinant
196 }
197}
198
199impl<R> From<Matrix<R>> for SparseMatrix<R>
200where
201 R: SemiRing<T: PartialEq>,
202{
203 fn from(mat: Matrix<R>) -> Self {
204 let mut nonzero = vec![];
205 for i in 0..mat.shape.0 {
206 for j in 0..mat.shape.1 {
207 let v = mat[(i, j)].clone();
208 if !R::is_zero(&v) {
209 nonzero.push((i, j, v));
210 }
211 }
212 }
213 Self {
214 shape: mat.shape,
215 nonzero,
216 }
217 }More examples
crates/competitive/src/algorithm/automata_learning.rs (line 493)
492 fn split_sample(&mut self, sample: &[usize]) -> Option<(Vec<usize>, Vec<usize>)> {
493 if self.prefixes.is_empty() && !F::is_zero(&self.automaton.behavior(sample.iter().cloned()))
494 {
495 return Some((vec![], sample.to_vec()));
496 }
497 let expected = self.automaton.behavior(sample.iter().cloned());
498 if expected == self.wfa.behavior(sample.iter().cloned()) {
499 return None;
500 }
501 let n = sample.len();
502 let dim = self.wfa.final_weights.shape.0;
503 let mut states: Vec<(Matrix<F>, usize)> = Vec::with_capacity(n + 1);
504 let mut v = self.wfa.final_weights.clone();
505 states.push((v.clone(), n));
506 for k in (0..n).rev() {
507 v = &self.wfa.transitions[sample[k]] * &v;
508 states.push((v.clone(), k));
509 }
510 states.reverse();
511 let split = states.partition_point(|(state, k)| {
512 (0..dim).any(|j| {
513 self.automaton.behavior(
514 self.prefixes[j]
515 .iter()
516 .cloned()
517 .chain(sample[*k..].iter().cloned()),
518 ) != state[j][0]
519 })
520 });
521 Some((sample[..split].to_vec(), sample[split..].to_vec()))
522 }
523 pub fn train_sample(&mut self, sample: &[usize]) -> bool {
524 let Some((prefix, suffix)) = self.split_sample(sample) else {
525 return false;
526 };
527 self.prefixes.push(prefix);
528 self.suffixes.push(suffix);
529 let n = self.inv_h.shape.0;
530 let prefix = &self.prefixes[n];
531 let suffix = &self.suffixes[n];
532 let u = Matrix::<F>::new_with((n, 1), |i, _| {
533 self.automaton.behavior(
534 self.prefixes[i]
535 .iter()
536 .cloned()
537 .chain(suffix.iter().cloned()),
538 )
539 });
540 let v = Matrix::<F>::new_with((1, n), |_, j| {
541 self.automaton.behavior(
542 prefix
543 .iter()
544 .cloned()
545 .chain(self.suffixes[j].iter().cloned()),
546 )
547 });
548 let w = Matrix::<F>::new_with((1, 1), |_, _| {
549 self.automaton
550 .behavior(prefix.iter().cloned().chain(suffix.iter().cloned()))
551 });
552 let t = &self.inv_h * &u;
553 let s = &v * &self.inv_h;
554 let d = F::inv(&(&w - &(&v * &t))[0][0]);
555 let dh = &t * &s;
556 for i in 0..n {
557 for j in 0..n {
558 F::add_assign(&mut self.inv_h[i][j], &F::mul(&dh[i][j], &d));
559 }
560 }
561 self.inv_h
562 .add_col_with(|i, _| F::neg(&F::mul(&t[i][0], &d)));
563 self.inv_h.add_row_with(|_, j| {
564 if j != n {
565 F::neg(&F::mul(&s[0][j], &d))
566 } else {
567 d.clone()
568 }
569 });
570
571 for (x, transition) in self.wfa.transitions.iter_mut().enumerate() {
572 let b = &(&self.nh[x] * &t) * &s;
573 for i in 0..n {
574 for j in 0..n {
575 F::add_assign(&mut transition[i][j], &F::mul(&b[i][j], &d));
576 }
577 }
578 }
579 for (x, nh) in self.nh.iter_mut().enumerate() {
580 nh.add_col_with(|i, j| {
581 self.automaton.behavior(
582 self.prefixes[i]
583 .iter()
584 .cloned()
585 .chain([x])
586 .chain(self.suffixes[j].iter().cloned()),
587 )
588 });
589 nh.add_row_with(|i, j| {
590 self.automaton.behavior(
591 self.prefixes[i]
592 .iter()
593 .cloned()
594 .chain([x])
595 .chain(self.suffixes[j].iter().cloned()),
596 )
597 });
598 }
599 self.wfa
600 .initial_weights
601 .add_col_with(|_, _| if n == 0 { F::one() } else { F::zero() });
602 self.wfa
603 .final_weights
604 .add_row_with(|_, _| self.automaton.behavior(prefix.iter().cloned()));
605 for (x, transition) in self.wfa.transitions.iter_mut().enumerate() {
606 transition.add_col_with(|_, _| F::zero());
607 transition.add_row_with(|_, _| F::zero());
608 for i in 0..=n {
609 for j in 0..=n {
610 if i == n || j == n {
611 for k in 0..=n {
612 if i != n && j != n && k != n {
613 continue;
614 }
615 F::add_assign(
616 &mut transition[i][k],
617 &F::mul(&self.nh[x][i][j], &self.inv_h[j][k]),
618 );
619 }
620 } else {
621 let k = n;
622 F::add_assign(
623 &mut transition[i][k],
624 &F::mul(&self.nh[x][i][j], &self.inv_h[j][k]),
625 );
626 }
627 }
628 }
629 }
630 true
631 }
632 pub fn train(&mut self, samples: impl IntoIterator<Item = Vec<usize>>) {
633 for sample in samples {
634 self.train_sample(&sample);
635 }
636 }
637 pub fn batch_train(&mut self, samples: impl IntoIterator<Item = Vec<usize>>) {
638 let mut prefix_set: HashSet<_> = self.prefixes.iter().cloned().collect();
639 let mut suffix_set: HashSet<_> = self.suffixes.iter().cloned().collect();
640 for sample in samples {
641 if prefix_set.insert(sample.to_vec()) {
642 self.prefixes.push(sample.to_vec());
643 }
644 if suffix_set.insert(sample.to_vec()) {
645 self.suffixes.push(sample);
646 }
647 }
648 let mut h = Matrix::<F>::new_with((self.prefixes.len(), self.suffixes.len()), |i, j| {
649 self.automaton.behavior(
650 self.prefixes[i]
651 .iter()
652 .cloned()
653 .chain(self.suffixes[j].iter().cloned()),
654 )
655 });
656 if !self.prefixes.is_empty() && !self.suffixes.is_empty() && F::is_zero(&h[0][0]) {
657 for j in 1..self.suffixes.len() {
658 if !F::is_zero(&h[0][j]) {
659 self.suffixes.swap(0, j);
660 for row in &mut h.data {
661 row.swap(0, j);
662 }
663 break;
664 }
665 }
666 }
667 let mut row_id: Vec<usize> = (0..h.shape.0).collect();
668 let mut pivots = vec![];
669 h.row_reduction_with(false, |r, p, c| {
670 row_id.swap(r, p);
671 pivots.push((row_id[r], c));
672 });
673 let mut new_prefixes = vec![];
674 let mut new_suffixes = vec![];
675 for (i, j) in pivots {
676 new_prefixes.push(self.prefixes[i].clone());
677 new_suffixes.push(self.suffixes[j].clone());
678 }
679 self.prefixes = new_prefixes;
680 self.suffixes = new_suffixes;
681 assert_eq!(self.prefixes.len(), self.suffixes.len());
682 let n = self.prefixes.len();
683 let h = Matrix::<F>::new_with((n, n), |i, j| {
684 self.automaton.behavior(
685 self.prefixes[i]
686 .iter()
687 .cloned()
688 .chain(self.suffixes[j].iter().cloned()),
689 )
690 });
691 self.inv_h = h.inverse().expect("Hankel matrix must be invertible");
692 self.wfa = WeightedFiniteAutomaton::<F> {
693 initial_weights: Matrix::new_with((1, n), |_, j| {
694 if self.prefixes[j].is_empty() {
695 F::one()
696 } else {
697 F::zero()
698 }
699 }),
700 transitions: (0..self.automaton.sigma())
701 .map(|x| {
702 &Matrix::new_with((n, n), |i, j| {
703 self.automaton.behavior(
704 self.prefixes[i]
705 .iter()
706 .cloned()
707 .chain([x])
708 .chain(self.suffixes[j].iter().cloned()),
709 )
710 }) * &self.inv_h
711 })
712 .collect(),
713 final_weights: Matrix::new_with((n, 1), |i, _| {
714 self.automaton.behavior(self.prefixes[i].iter().cloned())
715 }),
716 };
717 }crates/competitive/src/math/matrix.rs (line 174)
158 fn eliminate<const DETERMINANT: bool>(&mut self) -> (usize, R::T) {
159 let (n, m) = self.shape;
160 let mut rank = 0;
161 let mut determinant = R::one();
162 let mut negative = false;
163 for first in (0..m).step_by(64) {
164 if rank == n {
165 break;
166 }
167 let end = (first + 64).min(m);
168 let start = rank;
169 let mut pivots = Vec::new();
170 let mut panel = vec![vec![R::zero(); end - first]; end - first];
171 for col in first..end {
172 if panel[col - first][..col - first]
173 .iter()
174 .any(|x| !R::is_zero(x))
175 {
176 for row in &mut self.data[rank..] {
177 let value =
178 R::dot_product(&row[first..col], &panel[col - first][..col - first]);
179 R::sub_assign(&mut row[col], &value);
180 }
181 }
182 let Some(pivot) = (rank..n).find(|&i| !R::is_zero(&self[i][col])) else {
183 continue;
184 };
185 if pivot != rank {
186 self.data.swap(rank, pivot);
187 negative = !negative;
188 }
189 if DETERMINANT {
190 R::mul_assign(&mut determinant, &self[rank][col]);
191 }
192 let inv = R::inv(&self[rank][col]);
193 let row = &mut self.data[rank];
194 for c in col + 1..end {
195 let value = R::dot_product(&row[first..col], &panel[c - first][..col - first]);
196 R::sub_assign(&mut row[c], &value);
197 panel[c - first][col - first] = row[c].clone();
198 }
199 for row in &mut self.data[rank + 1..] {
200 R::mul_assign(&mut row[col], &inv);
201 }
202 pivots.push(col);
203 rank += 1;
204 if rank == n {
205 break;
206 }
207 }
208 for i in start..if rank - start < 32 { n } else { rank } {
209 let (upper, lower) = self.data.split_at_mut(i);
210 let row = &mut lower[0];
211 for (j, &col) in pivots[..(i - start).min(pivots.len())].iter().enumerate() {
212 if R::is_zero(&row[col]) {
213 continue;
214 }
215 let factor = R::neg(&row[col]);
216 R::add_scaled_assign(&mut row[end..], &upper[start + j][end..], &factor);
217 }
218 }
219 if rank < n
220 && end < m
221 && rank - start >= 32
222 && self.data[rank..]
223 .iter()
224 .any(|row| pivots.iter().any(|&col| !R::is_zero(&row[col])))
225 {
226 let lower = Self::new_with((n - rank, rank - start), |i, j| {
227 R::neg(&self[rank + i][pivots[j]])
228 });
229 let upper = Self::new_with((rank - start, m - end), |i, j| {
230 self[start + i][end + j].clone()
231 });
232 let update = &lower * &upper;
233 for (row, update) in self.data[rank..].iter_mut().zip(update.data) {
234 for (x, y) in row[end..].iter_mut().zip(update) {
235 R::add_assign(x, &y);
236 }
237 }
238 }
239 for (i, &col) in pivots.iter().enumerate() {
240 for row in &mut self.data[start + i + 1..] {
241 row[col] = R::zero();
242 }
243 }
244 if DETERMINANT && rank < end {
245 return (rank, R::zero());
246 }
247 }
248 if DETERMINANT && negative {
249 determinant = R::neg(&determinant);
250 }
251 (rank, determinant)
252 }
253
254 /// f: (row, pivot_row, col)
255 pub fn row_reduction_with<F>(&mut self, normalize: bool, mut f: F)
256 where
257 F: FnMut(usize, usize, usize),
258 {
259 let (n, m) = self.shape;
260 let mut c = 0;
261 for r in 0..n {
262 loop {
263 if c >= m {
264 return;
265 }
266 if let Some(pivot) = (r..n).find(|&p| !R::is_zero(&self[p][c])) {
267 f(r, pivot, c);
268 self.data.swap(r, pivot);
269 break;
270 };
271 c += 1;
272 }
273 let d = R::inv(&self[r][c]);
274 if normalize {
275 for value in &mut self[r][c..m] {
276 R::mul_assign(value, &d);
277 }
278 }
279 for i in (0..n).filter(|&i| i != r) {
280 let mut e = self[i][c].clone();
281 if !normalize {
282 R::mul_assign(&mut e, &d);
283 }
284 for j in c..m {
285 let e = R::mul(&e, &self[r][j]);
286 R::sub_assign(&mut self[i][j], &e);
287 }
288 }
289 c += 1;
290 }
291 }
292
293 pub fn row_reduction(&mut self, normalize: bool) {
294 self.row_reduction_with(normalize, |_, _, _| {});
295 }
296
297 pub fn rank(&mut self) -> usize {
298 self.eliminate::<false>().0
299 }
300
301 pub fn determinant(&mut self) -> R::T {
302 assert_eq!(self.shape.0, self.shape.1);
303 self.eliminate::<true>().1
304 }
305
306 pub fn solve_system_of_linear_equations(
307 &self,
308 b: &[R::T],
309 ) -> Option<SystemOfLinearEquationsSolution<R>> {
310 assert_eq!(self.shape.0, b.len());
311 let m = self.shape.1;
312 let mut a = Self::new_with((self.shape.0, m + 1), |i, j| {
313 if j == m {
314 b[i].clone()
315 } else {
316 self[i][j].clone()
317 }
318 });
319 let rank = a.eliminate::<false>().0;
320 let mut pivots = Vec::with_capacity(rank);
321 let mut b = Vec::with_capacity(rank);
322 for row in &a.data[..rank] {
323 let c = row.iter().position(|x| !R::is_zero(x)).unwrap();
324 if c == m {
325 return None;
326 }
327 pivots.push(c);
328 b.push(row[m].clone());
329 }
330
331 let mut free = Vec::with_capacity(m - rank);
332 let mut pivot = 0;
333 for c in 0..m {
334 if pivot < rank && pivots[pivot] == c {
335 pivot += 1;
336 } else {
337 free.push(c);
338 }
339 }
340 let mut coefficients: Vec<Vec<_>> = (0..rank)
341 .map(|i| free.iter().map(|&c| a[i][c].clone()).collect())
342 .collect();
343 for k in (0..rank).rev() {
344 let c = pivots[k];
345 let inv = R::inv(&a[k][c]);
346 R::mul_assign(&mut b[k], &inv);
347 let pivot_b = b[k].clone();
348 let (upper, lower) = coefficients.split_at_mut(k);
349 let pivot_coefficients = &mut lower[0];
350 for x in pivot_coefficients.iter_mut() {
351 R::mul_assign(x, &inv);
352 }
353 for ((row, value), coefficients) in a.data[..k].iter_mut().zip(&mut b[..k]).zip(upper) {
354 if R::is_zero(&row[c]) {
355 continue;
356 }
357 let factor = row[c].clone();
358 row[c] = R::zero();
359 R::sub_assign(value, &R::mul(&factor, &pivot_b));
360 R::add_scaled_assign(coefficients, pivot_coefficients, &R::neg(&factor));
361 }
362 }
363
364 let mut particular = vec![R::zero(); m];
365 for i in 0..rank {
366 particular[pivots[i]] = b[i].clone();
367 }
368 let mut basis = Vec::with_capacity(free.len());
369 for (j, &c) in free.iter().enumerate() {
370 let mut vector = vec![R::zero(); m];
371 vector[c] = R::one();
372 for i in 0..rank {
373 vector[pivots[i]] = R::neg(&coefficients[i][j]);
374 }
375 basis.push(vector);
376 }
377 Some(SystemOfLinearEquationsSolution { particular, basis })
378 }
379
380 pub fn inverse(&self) -> Option<Matrix<R>> {
381 assert_eq!(self.shape.0, self.shape.1);
382 let n = self.shape.0;
383 if n >= 64 {
384 let m = n / 2;
385 let a = Self::new_with((m, m), |i, j| self[i][j].clone());
386 if let Some(mut ai) = a.inverse() {
387 let b = Self::new_with((m, n - m), |i, j| self[i][j + m].clone());
388 let c = Self::new_with((n - m, m), |i, j| self[i + m][j].clone());
389 let mut d = Self::new_with((n - m, n - m), |i, j| self[i + m][j + m].clone());
390 let u = &ai * &b;
391 let v = &c * &ai;
392 d -= &v * &b;
393 let di = d.inverse()?;
394 let r = &u * &di;
395 let t = &di * &v;
396 ai += &r * &v;
397 let mut inverse = Self::zeros((n, n));
398 for i in 0..m {
399 inverse[i][..m].clone_from_slice(&ai[i]);
400 for (x, y) in inverse[i][m..].iter_mut().zip(&r[i]) {
401 *x = R::neg(y);
402 }
403 }
404 for i in m..n {
405 for (x, y) in inverse[i][..m].iter_mut().zip(&t[i - m]) {
406 *x = R::neg(y);
407 }
408 inverse[i][m..].clone_from_slice(&di[i - m]);
409 }
410 return Some(inverse);
411 }
412 }
413 let mut a = self.clone();
414 let mut inverse = Self::eye((n, n));
415 let mut ranges: Vec<_> = (0..n).map(|i| (i, i + 1)).collect();
416 for r in 0..n {
417 let pivot = (r..n).find(|&i| !R::is_zero(&a[i][r]))?;
418 a.data.swap(r, pivot);
419 inverse.data.swap(r, pivot);
420 ranges.swap(r, pivot);
421
422 let d = R::inv(&a[r][r]);
423 for x in &mut a[r][r..] {
424 R::mul_assign(x, &d);
425 }
426 let (left, right) = ranges[r];
427 for x in &mut inverse[r][left..right] {
428 R::mul_assign(x, &d);
429 }
430
431 let (a_upper, a_lower) = a.data.split_at_mut(r + 1);
432 let pivot_a = &a_upper[r];
433 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r + 1);
434 let pivot_inverse = &inverse_upper[r];
435 let (ranges_upper, ranges_lower) = ranges.split_at_mut(r + 1);
436 let (left, right) = ranges_upper[r];
437 for ((a, inverse), range) in a_lower.iter_mut().zip(inverse_lower).zip(ranges_lower) {
438 if R::is_zero(&a[r]) {
439 continue;
440 }
441 let e = a[r].clone();
442 a[r] = R::zero();
443 R::add_scaled_assign(&mut a[(r + 1)..], &pivot_a[(r + 1)..], &R::neg(&e));
444 R::add_scaled_assign(
445 &mut inverse[left..right],
446 &pivot_inverse[left..right],
447 &R::neg(&e),
448 );
449 range.0 = range.0.min(left);
450 range.1 = range.1.max(right);
451 }
452 }
453 for r in (0..n).rev() {
454 let (left, right) = ranges[r];
455 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r);
456 let pivot_inverse = &inverse_lower[0];
457 let (ranges_upper, _) = ranges.split_at_mut(r);
458 for ((a, inverse), range) in a.data[..r].iter_mut().zip(inverse_upper).zip(ranges_upper)
459 {
460 if R::is_zero(&a[r]) {
461 continue;
462 }
463 let e = a[r].clone();
464 a[r] = R::zero();
465 R::add_scaled_assign(
466 &mut inverse[left..right],
467 &pivot_inverse[left..right],
468 &R::neg(&e),
469 );
470 range.0 = range.0.min(left);
471 range.1 = range.1.max(right);
472 }
473 }
474 Some(inverse)
475 }
476
477 pub fn characteristic_polynomial(&mut self) -> Vec<R::T> {
478 let n = self.shape.0;
479 if n == 0 {
480 return vec![R::one()];
481 }
482 assert!(self.data.iter().all(|a| a.len() == n));
483 for j in 0..(n - 1) {
484 if let Some(x) = ((j + 1)..n).find(|&x| !R::is_zero(&self[x][j])) {
485 self.data.swap(j + 1, x);
486 self.data.iter_mut().for_each(|a| a.swap(j + 1, x));
487 let inv = R::inv(&self[j + 1][j]);
488 let mut v = vec![];
489 let src = std::mem::take(&mut self[j + 1]);
490 for a in self.data[(j + 2)..].iter_mut() {
491 let mul = R::mul(&a[j], &inv);
492 R::add_scaled_assign(&mut a[j..], &src[j..], &R::neg(&mul));
493 v.push(mul);
494 }
495 self[j + 1] = src;
496 for a in self.data.iter_mut() {
497 let v = R::dot_product(&a[(j + 2)..], &v);
498 R::add_assign(&mut a[j + 1], &v);
499 }
500 }
501 }
502 // dp[k][j - k] stores [x^k] det(xI - A[..j, ..j]).
503 let mut dp: Vec<Vec<R::T>> = (0..=n).map(|i| Vec::with_capacity(n + 1 - i)).collect();
504 dp[0].push(R::one());
505 for i in 0..n {
506 let mut c = vec![R::zero(); i + 1];
507 c[i] = R::neg(&self[i][i]);
508 let mut mul = R::one();
509 for j in (0..i).rev() {
510 mul = R::mul(&mul, &self[j + 1][j]);
511 c[j] = R::neg(&R::mul(&mul, &self[j][i]));
512 }
513 for k in (0..=i).rev() {
514 let mut value = R::dot_product(&dp[k], &c[k..]);
515 if k > 0 {
516 R::add_assign(&mut value, dp[k - 1].last().unwrap());
517 }
518 dp[k].push(value);
519 }
520 dp[i + 1].push(R::one());
521 }
522 dp.into_iter().map(|mut c| c.pop().unwrap()).collect()
523 }Sourcefn one() -> Self::T
fn one() -> Self::T
multiplicative identity: $1$
Examples found in repository?
More examples
crates/competitive/src/string/rolling_hash.rs (line 500)
497 fn new(base: R::T) -> Self {
498 Self {
499 base,
500 pow: vec![R::one()],
501 }
502 }
503 fn ensure_pow(&mut self, len: usize) {
504 if self.pow.len() <= len {
505 self.pow.reserve(len - self.pow.len() + 1);
506 if self.pow.is_empty() {
507 self.pow.push(R::one());
508 }
509 for _ in 0..=len - self.pow.len() {
510 self.pow.push(R::mul(self.pow.last().unwrap(), &self.base));
511 }
512 }
513 }crates/competitive/src/math/floor_sum.rs (line 175)
173 fn to_x() -> FloorSumData<R, X, Y> {
174 let mut dp = array![array![R::zero(); Y]; X];
175 dp[0][0] = R::one();
176 FloorSumData {
177 dp,
178 dx: R::one(),
179 dy: R::zero(),
180 _marker: PhantomData,
181 }
182 }
183 fn to_y() -> FloorSumData<R, X, Y> {
184 FloorSumData {
185 dp: array![array![R::zero(); Y]; X],
186 dx: R::zero(),
187 dy: R::one(),
188 _marker: PhantomData,
189 }
190 }
191}
192
193impl<R, const X: usize, const Y: usize> FloorSum<R, X, Y>
194where
195 R: Ring<Additive: Invertible>,
196{
197 fn offset(x: i64, y: i64) -> FloorSumData<R, X, Y> {
198 FloorSumData {
199 dp: array![array![R::zero(); Y]; X],
200 dx: R::Additive::signed_pow(R::one(), x),
201 dy: R::Additive::signed_pow(R::one(), y),
202 _marker: PhantomData,
203 }
204 }
205}
206
207impl<R, const X: usize, const Y: usize> Magma for FloorSum<R, X, Y>
208where
209 R: SemiRing,
210{
211 type T = FloorSumData<R, X, Y>;
212
213 fn operate(a: &Self::T, b: &Self::T) -> Self::T {
214 let mut a = a.clone();
215 let mut b = b.clone();
216 let mut pow_x = array![R::zero(); X];
217 let mut pow_y = array![R::zero(); Y];
218 pow_x[0] = R::one();
219 pow_y[0] = R::one();
220 for i in 1..X {
221 pow_x[i] = R::mul(&pow_x[i - 1], &a.dx);
222 }
223 for j in 1..Y {
224 pow_y[j] = R::mul(&pow_y[j - 1], &a.dy);
225 }
226 macro_rules! go {
227 ($N:ident) => {
228 let mut comb = array![array![R::zero(); $N]; $N];
229 comb[0][0] = R::one();
230 let mut i = 0;
231 while i + 1 < $N {
232 let mut j = 0;
233 while j <= i {
234 let x = comb[i][j].clone();
235 R::add_assign(&mut comb[i + 1][j], &x);
236 R::add_assign(&mut comb[i + 1][j + 1], &x);
237 j += 1;
238 }
239 i += 1;
240 }
241 for i in 0..X {
242 for j in (0..Y).rev() {
243 for k in j + 1..Y {
244 let mut x = b.dp[i][j].clone();
245 R::mul_assign(&mut x, &comb[k][j]);
246 R::mul_assign(&mut x, &pow_y[k - j]);
247 R::add_assign(&mut b.dp[i][k], &x);
248 }
249 }
250 }
251 for j in 0..Y {
252 for i in (0..X).rev() {
253 for k in i..X {
254 let mut x = b.dp[i][j].clone();
255 R::mul_assign(&mut x, &comb[k][i]);
256 R::mul_assign(&mut x, &pow_x[k - i]);
257 R::add_assign(&mut a.dp[k][j], &x);
258 }
259 }
260 }
261 };
262 }
263 if X <= Y {
264 go!(Y);
265 } else {
266 go!(X);
267 }
268 R::add_assign(&mut a.dx, &b.dx);
269 R::add_assign(&mut a.dy, &b.dy);
270 a
271 }
272}
273
274impl<R, const X: usize, const Y: usize> Unital for FloorSum<R, X, Y>
275where
276 R: SemiRing,
277{
278 fn unit() -> Self::T {
279 FloorSumData {
280 dp: array![array![R::zero(); Y]; X],
281 dx: R::zero(),
282 dy: R::zero(),
283 _marker: PhantomData,
284 }
285 }
286}
287
288impl<R, const X: usize, const Y: usize> Associative for FloorSum<R, X, Y> where R: SemiRing {}
289
290fn floor_monoid_product<M>(
291 mut x: M::T,
292 mut y: M::T,
293 mut n: u64,
294 mut a: u64,
295 mut b: u64,
296 mut m: u64,
297) -> M::T
298where
299 M: Monoid,
300{
301 let mut c = (a * n + b) / m;
302 let mut pre = M::unit();
303 let mut suf = M::unit();
304 loop {
305 let (p, q) = (a / m, b / m);
306 a %= m;
307 b %= m;
308 x = M::operate(&x, &M::pow(y.clone(), p));
309 pre = M::operate(&pre, &M::pow(y.clone(), q));
310 c -= p * n + q;
311 if c == 0 {
312 break;
313 }
314 let d = (m * c - b - 1) / a + 1;
315 suf = M::operate(&y, &M::operate(&M::pow(x.clone(), n - d), &suf));
316 b = m - b - 1 + a;
317 n = c - 1;
318 c = d;
319 swap(&mut m, &mut a);
320 swap(&mut x, &mut y);
321 }
322 x = M::pow(x.clone(), n);
323 M::operate(&M::operate(&pre, &x), &suf)
324}
325
326/// $$\sum_{i=0}^{n-1}i^X\left\lfloor\frac{a\times i+b}{m}\right\rfloor^Y$$
327pub fn floor_sum_polynomial<T, const X: usize, const Y: usize>(
328 n: u64,
329 a: u64,
330 b: u64,
331 m: u64,
332) -> [[T; Y]; X]
333where
334 T: DotProduct + One,
335{
336 debug_assert!(a == 0 || n < (u64::MAX - b) / a);
337 floor_monoid_product::<FloorSum<AddMulOperation<T>, X, Y>>(
338 FloorSum::<AddMulOperation<T>, X, Y>::to_x(),
339 FloorSum::<AddMulOperation<T>, X, Y>::to_y(),
340 n,
341 a,
342 b,
343 m,
344 )
345 .dp
346}
347
348/// $$\sum_{i=l}^{r-1}i^X\left\lfloor\frac{a\times i+b}{m}\right\rfloor^Y$$
349pub fn floor_sum_polynomial_i64<T, const X: usize, const Y: usize>(
350 l: i64,
351 r: i64,
352 a: i64,
353 b: i64,
354 m: u64,
355) -> [[T; Y]; X]
356where
357 T: DotProduct + One,
358 AddMulOperation<T>: SemiRing<T = T, Additive: Invertible>,
359{
360 assert!(l <= r);
361 assert!(m > 0);
362
363 if a < 0 {
364 let mut ans = floor_sum_polynomial_i64::<T, X, Y>(-r + 1, -l + 1, -a, b, m);
365 for ans in ans.iter_mut().skip(1).step_by(2) {
366 for ans in ans.iter_mut() {
367 *ans = AddMulOperation::<T>::neg(ans);
368 }
369 }
370 return ans;
371 }
372
373 let add_x = l;
374 let n = (r - l) as u64;
375 let b = a * add_x + b;
376
377 let add_y = b.div_euclid(m as i64);
378 let b = b.rem_euclid(m as i64);
379 assert!(a >= 0);
380 assert!(b >= 0);
381 let data = floor_monoid_product::<FloorSum<AddMulOperation<T>, X, Y>>(
382 FloorSum::<AddMulOperation<T>, X, Y>::to_x(),
383 FloorSum::<AddMulOperation<T>, X, Y>::to_y(),
384 n,
385 a as u64,
386 b as u64,
387 m,
388 );
389
390 let offset = FloorSum::<AddMulOperation<T>, X, Y>::offset(add_x, add_y);
391 FloorSum::<AddMulOperation<T>, X, Y>::operate(&offset, &data).dp
392}
393
394#[derive(Debug)]
395struct FloorPowerSum<R>
396where
397 R: SemiRing,
398{
399 x: R::T,
400 sum: R::T,
401}
402
403impl<R> Clone for FloorPowerSum<R>
404where
405 R: SemiRing,
406{
407 fn clone(&self) -> Self {
408 Self {
409 x: self.x.clone(),
410 sum: self.sum.clone(),
411 }
412 }
413}
414
415impl<R> FloorPowerSum<R>
416where
417 R: SemiRing,
418{
419 fn to_x(x: R::T) -> Self {
420 Self { x, sum: R::one() }
421 }
422 fn to_y(y: R::T) -> Self {
423 Self {
424 x: y,
425 sum: R::zero(),
426 }
427 }
428}
429
430impl<R> Magma for FloorPowerSum<R>
431where
432 R: SemiRing,
433{
434 type T = Self;
435
436 fn operate(a: &Self::T, b: &Self::T) -> Self::T {
437 Self {
438 x: R::mul(&a.x, &b.x),
439 sum: R::add(&a.sum, &R::mul(&a.x, &b.sum)),
440 }
441 }
442}
443
444impl<R> Unital for FloorPowerSum<R>
445where
446 R: SemiRing,
447{
448 fn unit() -> Self::T {
449 Self {
450 x: R::one(),
451 sum: R::zero(),
452 }
453 }crates/competitive/src/math/matrix.rs (line 98)
95 pub fn eye(shape: (usize, usize)) -> Self {
96 let mut data = vec![vec![R::zero(); shape.1]; shape.0];
97 for (i, d) in data.iter_mut().enumerate().take(shape.1) {
98 d[i] = R::one();
99 }
100 Self {
101 shape,
102 data,
103 _marker: PhantomData,
104 }
105 }
106
107 pub fn transpose(&self) -> Self {
108 Self::new_with((self.shape.1, self.shape.0), |i, j| self[j][i].clone())
109 }
110
111 pub fn map<S, F>(&self, mut f: F) -> Matrix<S>
112 where
113 S: SemiRing,
114 F: FnMut(&R::T) -> S::T,
115 {
116 Matrix::<S>::new_with(self.shape, |i, j| f(&self[i][j]))
117 }
118
119 pub fn add_row_with(&mut self, mut f: impl FnMut(usize, usize) -> R::T) {
120 self.data
121 .push((0..self.shape.1).map(|j| f(self.shape.0, j)).collect());
122 self.shape.0 += 1;
123 }
124
125 pub fn add_col_with(&mut self, mut f: impl FnMut(usize, usize) -> R::T) {
126 for i in 0..self.shape.0 {
127 self.data[i].push(f(i, self.shape.1));
128 }
129 self.shape.1 += 1;
130 }
131
132 pub fn pairwise_assign<F>(&mut self, other: &Self, mut f: F)
133 where
134 F: FnMut(&mut R::T, &R::T),
135 {
136 assert_eq!(self.shape, other.shape);
137 for i in 0..self.shape.0 {
138 for j in 0..self.shape.1 {
139 f(&mut self[i][j], &other[i][j]);
140 }
141 }
142 }
143}
144
145#[derive(Debug)]
146pub struct SystemOfLinearEquationsSolution<R>
147where
148 R: Field<Additive: Invertible, Multiplicative: Invertible>,
149{
150 pub particular: Vec<R::T>,
151 pub basis: Vec<Vec<R::T>>,
152}
153
154impl<R> Matrix<R>
155where
156 R: Field<T: PartialEq, Additive: Invertible, Multiplicative: Invertible>,
157{
158 fn eliminate<const DETERMINANT: bool>(&mut self) -> (usize, R::T) {
159 let (n, m) = self.shape;
160 let mut rank = 0;
161 let mut determinant = R::one();
162 let mut negative = false;
163 for first in (0..m).step_by(64) {
164 if rank == n {
165 break;
166 }
167 let end = (first + 64).min(m);
168 let start = rank;
169 let mut pivots = Vec::new();
170 let mut panel = vec![vec![R::zero(); end - first]; end - first];
171 for col in first..end {
172 if panel[col - first][..col - first]
173 .iter()
174 .any(|x| !R::is_zero(x))
175 {
176 for row in &mut self.data[rank..] {
177 let value =
178 R::dot_product(&row[first..col], &panel[col - first][..col - first]);
179 R::sub_assign(&mut row[col], &value);
180 }
181 }
182 let Some(pivot) = (rank..n).find(|&i| !R::is_zero(&self[i][col])) else {
183 continue;
184 };
185 if pivot != rank {
186 self.data.swap(rank, pivot);
187 negative = !negative;
188 }
189 if DETERMINANT {
190 R::mul_assign(&mut determinant, &self[rank][col]);
191 }
192 let inv = R::inv(&self[rank][col]);
193 let row = &mut self.data[rank];
194 for c in col + 1..end {
195 let value = R::dot_product(&row[first..col], &panel[c - first][..col - first]);
196 R::sub_assign(&mut row[c], &value);
197 panel[c - first][col - first] = row[c].clone();
198 }
199 for row in &mut self.data[rank + 1..] {
200 R::mul_assign(&mut row[col], &inv);
201 }
202 pivots.push(col);
203 rank += 1;
204 if rank == n {
205 break;
206 }
207 }
208 for i in start..if rank - start < 32 { n } else { rank } {
209 let (upper, lower) = self.data.split_at_mut(i);
210 let row = &mut lower[0];
211 for (j, &col) in pivots[..(i - start).min(pivots.len())].iter().enumerate() {
212 if R::is_zero(&row[col]) {
213 continue;
214 }
215 let factor = R::neg(&row[col]);
216 R::add_scaled_assign(&mut row[end..], &upper[start + j][end..], &factor);
217 }
218 }
219 if rank < n
220 && end < m
221 && rank - start >= 32
222 && self.data[rank..]
223 .iter()
224 .any(|row| pivots.iter().any(|&col| !R::is_zero(&row[col])))
225 {
226 let lower = Self::new_with((n - rank, rank - start), |i, j| {
227 R::neg(&self[rank + i][pivots[j]])
228 });
229 let upper = Self::new_with((rank - start, m - end), |i, j| {
230 self[start + i][end + j].clone()
231 });
232 let update = &lower * &upper;
233 for (row, update) in self.data[rank..].iter_mut().zip(update.data) {
234 for (x, y) in row[end..].iter_mut().zip(update) {
235 R::add_assign(x, &y);
236 }
237 }
238 }
239 for (i, &col) in pivots.iter().enumerate() {
240 for row in &mut self.data[start + i + 1..] {
241 row[col] = R::zero();
242 }
243 }
244 if DETERMINANT && rank < end {
245 return (rank, R::zero());
246 }
247 }
248 if DETERMINANT && negative {
249 determinant = R::neg(&determinant);
250 }
251 (rank, determinant)
252 }
253
254 /// f: (row, pivot_row, col)
255 pub fn row_reduction_with<F>(&mut self, normalize: bool, mut f: F)
256 where
257 F: FnMut(usize, usize, usize),
258 {
259 let (n, m) = self.shape;
260 let mut c = 0;
261 for r in 0..n {
262 loop {
263 if c >= m {
264 return;
265 }
266 if let Some(pivot) = (r..n).find(|&p| !R::is_zero(&self[p][c])) {
267 f(r, pivot, c);
268 self.data.swap(r, pivot);
269 break;
270 };
271 c += 1;
272 }
273 let d = R::inv(&self[r][c]);
274 if normalize {
275 for value in &mut self[r][c..m] {
276 R::mul_assign(value, &d);
277 }
278 }
279 for i in (0..n).filter(|&i| i != r) {
280 let mut e = self[i][c].clone();
281 if !normalize {
282 R::mul_assign(&mut e, &d);
283 }
284 for j in c..m {
285 let e = R::mul(&e, &self[r][j]);
286 R::sub_assign(&mut self[i][j], &e);
287 }
288 }
289 c += 1;
290 }
291 }
292
293 pub fn row_reduction(&mut self, normalize: bool) {
294 self.row_reduction_with(normalize, |_, _, _| {});
295 }
296
297 pub fn rank(&mut self) -> usize {
298 self.eliminate::<false>().0
299 }
300
301 pub fn determinant(&mut self) -> R::T {
302 assert_eq!(self.shape.0, self.shape.1);
303 self.eliminate::<true>().1
304 }
305
306 pub fn solve_system_of_linear_equations(
307 &self,
308 b: &[R::T],
309 ) -> Option<SystemOfLinearEquationsSolution<R>> {
310 assert_eq!(self.shape.0, b.len());
311 let m = self.shape.1;
312 let mut a = Self::new_with((self.shape.0, m + 1), |i, j| {
313 if j == m {
314 b[i].clone()
315 } else {
316 self[i][j].clone()
317 }
318 });
319 let rank = a.eliminate::<false>().0;
320 let mut pivots = Vec::with_capacity(rank);
321 let mut b = Vec::with_capacity(rank);
322 for row in &a.data[..rank] {
323 let c = row.iter().position(|x| !R::is_zero(x)).unwrap();
324 if c == m {
325 return None;
326 }
327 pivots.push(c);
328 b.push(row[m].clone());
329 }
330
331 let mut free = Vec::with_capacity(m - rank);
332 let mut pivot = 0;
333 for c in 0..m {
334 if pivot < rank && pivots[pivot] == c {
335 pivot += 1;
336 } else {
337 free.push(c);
338 }
339 }
340 let mut coefficients: Vec<Vec<_>> = (0..rank)
341 .map(|i| free.iter().map(|&c| a[i][c].clone()).collect())
342 .collect();
343 for k in (0..rank).rev() {
344 let c = pivots[k];
345 let inv = R::inv(&a[k][c]);
346 R::mul_assign(&mut b[k], &inv);
347 let pivot_b = b[k].clone();
348 let (upper, lower) = coefficients.split_at_mut(k);
349 let pivot_coefficients = &mut lower[0];
350 for x in pivot_coefficients.iter_mut() {
351 R::mul_assign(x, &inv);
352 }
353 for ((row, value), coefficients) in a.data[..k].iter_mut().zip(&mut b[..k]).zip(upper) {
354 if R::is_zero(&row[c]) {
355 continue;
356 }
357 let factor = row[c].clone();
358 row[c] = R::zero();
359 R::sub_assign(value, &R::mul(&factor, &pivot_b));
360 R::add_scaled_assign(coefficients, pivot_coefficients, &R::neg(&factor));
361 }
362 }
363
364 let mut particular = vec![R::zero(); m];
365 for i in 0..rank {
366 particular[pivots[i]] = b[i].clone();
367 }
368 let mut basis = Vec::with_capacity(free.len());
369 for (j, &c) in free.iter().enumerate() {
370 let mut vector = vec![R::zero(); m];
371 vector[c] = R::one();
372 for i in 0..rank {
373 vector[pivots[i]] = R::neg(&coefficients[i][j]);
374 }
375 basis.push(vector);
376 }
377 Some(SystemOfLinearEquationsSolution { particular, basis })
378 }
379
380 pub fn inverse(&self) -> Option<Matrix<R>> {
381 assert_eq!(self.shape.0, self.shape.1);
382 let n = self.shape.0;
383 if n >= 64 {
384 let m = n / 2;
385 let a = Self::new_with((m, m), |i, j| self[i][j].clone());
386 if let Some(mut ai) = a.inverse() {
387 let b = Self::new_with((m, n - m), |i, j| self[i][j + m].clone());
388 let c = Self::new_with((n - m, m), |i, j| self[i + m][j].clone());
389 let mut d = Self::new_with((n - m, n - m), |i, j| self[i + m][j + m].clone());
390 let u = &ai * &b;
391 let v = &c * &ai;
392 d -= &v * &b;
393 let di = d.inverse()?;
394 let r = &u * &di;
395 let t = &di * &v;
396 ai += &r * &v;
397 let mut inverse = Self::zeros((n, n));
398 for i in 0..m {
399 inverse[i][..m].clone_from_slice(&ai[i]);
400 for (x, y) in inverse[i][m..].iter_mut().zip(&r[i]) {
401 *x = R::neg(y);
402 }
403 }
404 for i in m..n {
405 for (x, y) in inverse[i][..m].iter_mut().zip(&t[i - m]) {
406 *x = R::neg(y);
407 }
408 inverse[i][m..].clone_from_slice(&di[i - m]);
409 }
410 return Some(inverse);
411 }
412 }
413 let mut a = self.clone();
414 let mut inverse = Self::eye((n, n));
415 let mut ranges: Vec<_> = (0..n).map(|i| (i, i + 1)).collect();
416 for r in 0..n {
417 let pivot = (r..n).find(|&i| !R::is_zero(&a[i][r]))?;
418 a.data.swap(r, pivot);
419 inverse.data.swap(r, pivot);
420 ranges.swap(r, pivot);
421
422 let d = R::inv(&a[r][r]);
423 for x in &mut a[r][r..] {
424 R::mul_assign(x, &d);
425 }
426 let (left, right) = ranges[r];
427 for x in &mut inverse[r][left..right] {
428 R::mul_assign(x, &d);
429 }
430
431 let (a_upper, a_lower) = a.data.split_at_mut(r + 1);
432 let pivot_a = &a_upper[r];
433 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r + 1);
434 let pivot_inverse = &inverse_upper[r];
435 let (ranges_upper, ranges_lower) = ranges.split_at_mut(r + 1);
436 let (left, right) = ranges_upper[r];
437 for ((a, inverse), range) in a_lower.iter_mut().zip(inverse_lower).zip(ranges_lower) {
438 if R::is_zero(&a[r]) {
439 continue;
440 }
441 let e = a[r].clone();
442 a[r] = R::zero();
443 R::add_scaled_assign(&mut a[(r + 1)..], &pivot_a[(r + 1)..], &R::neg(&e));
444 R::add_scaled_assign(
445 &mut inverse[left..right],
446 &pivot_inverse[left..right],
447 &R::neg(&e),
448 );
449 range.0 = range.0.min(left);
450 range.1 = range.1.max(right);
451 }
452 }
453 for r in (0..n).rev() {
454 let (left, right) = ranges[r];
455 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r);
456 let pivot_inverse = &inverse_lower[0];
457 let (ranges_upper, _) = ranges.split_at_mut(r);
458 for ((a, inverse), range) in a.data[..r].iter_mut().zip(inverse_upper).zip(ranges_upper)
459 {
460 if R::is_zero(&a[r]) {
461 continue;
462 }
463 let e = a[r].clone();
464 a[r] = R::zero();
465 R::add_scaled_assign(
466 &mut inverse[left..right],
467 &pivot_inverse[left..right],
468 &R::neg(&e),
469 );
470 range.0 = range.0.min(left);
471 range.1 = range.1.max(right);
472 }
473 }
474 Some(inverse)
475 }
476
477 pub fn characteristic_polynomial(&mut self) -> Vec<R::T> {
478 let n = self.shape.0;
479 if n == 0 {
480 return vec![R::one()];
481 }
482 assert!(self.data.iter().all(|a| a.len() == n));
483 for j in 0..(n - 1) {
484 if let Some(x) = ((j + 1)..n).find(|&x| !R::is_zero(&self[x][j])) {
485 self.data.swap(j + 1, x);
486 self.data.iter_mut().for_each(|a| a.swap(j + 1, x));
487 let inv = R::inv(&self[j + 1][j]);
488 let mut v = vec![];
489 let src = std::mem::take(&mut self[j + 1]);
490 for a in self.data[(j + 2)..].iter_mut() {
491 let mul = R::mul(&a[j], &inv);
492 R::add_scaled_assign(&mut a[j..], &src[j..], &R::neg(&mul));
493 v.push(mul);
494 }
495 self[j + 1] = src;
496 for a in self.data.iter_mut() {
497 let v = R::dot_product(&a[(j + 2)..], &v);
498 R::add_assign(&mut a[j + 1], &v);
499 }
500 }
501 }
502 // dp[k][j - k] stores [x^k] det(xI - A[..j, ..j]).
503 let mut dp: Vec<Vec<R::T>> = (0..=n).map(|i| Vec::with_capacity(n + 1 - i)).collect();
504 dp[0].push(R::one());
505 for i in 0..n {
506 let mut c = vec![R::zero(); i + 1];
507 c[i] = R::neg(&self[i][i]);
508 let mut mul = R::one();
509 for j in (0..i).rev() {
510 mul = R::mul(&mul, &self[j + 1][j]);
511 c[j] = R::neg(&R::mul(&mul, &self[j][i]));
512 }
513 for k in (0..=i).rev() {
514 let mut value = R::dot_product(&dp[k], &c[k..]);
515 if k > 0 {
516 R::add_assign(&mut value, dp[k - 1].last().unwrap());
517 }
518 dp[k].push(value);
519 }
520 dp[i + 1].push(R::one());
521 }
522 dp.into_iter().map(|mut c| c.pop().unwrap()).collect()
523 }crates/competitive/src/math/black_box_matrix.rs (line 133)
97 pub fn determinant(&self) -> R::T {
98 assert_eq!(self.shape.0, self.shape.1);
99 let n = self.shape.0;
100 let mut columns = vec![Vec::<(usize, R::T)>::new(); n];
101 for &(i, j, ref value) in &self.nonzero {
102 columns[j].push((i, value.clone()));
103 }
104 let mut degrees = vec![0; n];
105 for column in &mut columns {
106 column.sort_unstable_by_key(|&(i, _)| i);
107 let mut merged: Vec<(usize, R::T)> = Vec::with_capacity(column.len());
108 for (i, value) in column.drain(..) {
109 if let Some((_, x)) = merged.last_mut().filter(|(last, _)| *last == i) {
110 R::add_assign(x, &value);
111 } else {
112 merged.push((i, value));
113 }
114 }
115 merged.retain(|(i, value)| {
116 if R::is_zero(value) {
117 false
118 } else {
119 degrees[*i] += 1;
120 true
121 }
122 });
123 *column = merged;
124 }
125 let mut order: Vec<_> = (0..n).collect();
126 order.sort_unstable_by_key(|&j| columns[j].len());
127 let mut lower: Vec<Vec<(usize, R::T)>> = Vec::with_capacity(n);
128 let mut pivots: Vec<Option<usize>> = vec![None; n];
129 let mut x = vec![R::zero(); n];
130 let mut seen = vec![0; n];
131 let mut stack = Vec::new();
132 let mut support = Vec::new();
133 let mut determinant = R::one();
134 for (k, &j) in order.iter().enumerate() {
135 support.clear();
136 for &(i, _) in &columns[j] {
137 if seen[i] == k + 1 {
138 continue;
139 }
140 seen[i] = k + 1;
141 x[i] = R::zero();
142 stack.push((i, 0));
143 while let Some((i, next)) = stack.last_mut() {
144 if let Some(pivot) = pivots[*i]
145 && *next < lower[pivot].len()
146 {
147 let row = lower[pivot][*next].0;
148 *next += 1;
149 if seen[row] != k + 1 {
150 seen[row] = k + 1;
151 x[row] = R::zero();
152 stack.push((row, 0));
153 }
154 continue;
155 }
156 support.push(*i);
157 stack.pop();
158 }
159 }
160 for &(i, ref value) in &columns[j] {
161 x[i] = value.clone();
162 }
163 let mut pivot = None;
164 for &i in support.iter().rev() {
165 if let Some(p) = pivots[i] {
166 let factor = R::neg(&x[i]);
167 for &(row, ref value) in &lower[p] {
168 R::add_assign(&mut x[row], &R::mul(&factor, value));
169 }
170 } else if !R::is_zero(&x[i]) && pivot.is_none_or(|p| degrees[i] < degrees[p]) {
171 pivot = Some(i);
172 }
173 }
174 let Some(pivot) = pivot else { return R::zero() };
175 R::mul_assign(&mut determinant, &x[pivot]);
176 let inv = R::inv(&x[pivot]);
177 pivots[pivot] = Some(k);
178 lower.push(
179 support
180 .iter()
181 .filter(|&&i| pivots[i].is_none() && !R::is_zero(&x[i]))
182 .map(|&i| (i, R::mul(&x[i], &inv)))
183 .collect(),
184 );
185 }
186 for mut permutation in [order, pivots.into_iter().map(Option::unwrap).collect()] {
187 for i in 0..n {
188 while permutation[i] != i {
189 let j = permutation[i];
190 permutation.swap(i, j);
191 determinant = R::neg(&determinant);
192 }
193 }
194 }
195 determinant
196 }Additional examples can be found in:
fn is_one(x: &Self::T) -> bool
Sourcefn add(x: &Self::T, y: &Self::T) -> Self::T
fn add(x: &Self::T, y: &Self::T) -> Self::T
additive operaion: $+$
Examples found in repository?
More examples
crates/competitive/src/graph/shortest_path.rs (line 110)
107 fn add_assign(x: &mut Self::T, y: &Self::T) -> bool {
108 match x.0.cmp(&y.0) {
109 Ordering::Equal => {
110 x.1 = S::add(&x.1, &y.1);
111 false
112 }
113 Ordering::Greater => {
114 *x = y.clone();
115 true
116 }
117 _ => false,
118 }
119 }crates/competitive/src/algorithm/esper.rs (line 152)
140 pub fn solve(&self, input: Input) -> R::T {
141 let coeff = self
142 .data
143 .get(&(self.class)(&input))
144 .expect("unrecognized class")
145 .as_ref()
146 .expect("failed to solve");
147 let feature = (self.feature)(&input);
148 feature
149 .into_iter()
150 .zip(coeff)
151 .map(|(x, y)| R::mul(&x, y))
152 .fold(R::zero(), |x, y| R::add(&x, &y))
153 }crates/competitive/src/math/matrix.rs (line 702)
679fn strassen_rec<R: Ring>(
680 a: &[R::T],
681 b: &[R::T],
682 c: &mut [R::T],
683 shape: (usize, usize, usize),
684 stride_a: usize,
685 stride_b: usize,
686) {
687 let (n, m, p) = shape;
688 fn add_block<R: Ring>(
689 a: &[R::T],
690 b: &[R::T],
691 out: &mut [R::T],
692 n: usize,
693 stride_a: usize,
694 stride_b: usize,
695 ) {
696 for ((a, b), c) in a
697 .chunks(stride_a)
698 .zip(b.chunks(stride_b))
699 .zip(out.chunks_exact_mut(n))
700 {
701 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
702 *c = R::add(a, b);
703 }
704 }
705 }
706
707 fn sub_block<R: Ring>(
708 a: &[R::T],
709 b: &[R::T],
710 out: &mut [R::T],
711 n: usize,
712 stride_a: usize,
713 stride_b: usize,
714 ) {
715 for ((a, b), c) in a
716 .chunks(stride_a)
717 .zip(b.chunks(stride_b))
718 .zip(out.chunks_exact_mut(n))
719 {
720 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
721 *c = R::sub(a, b);
722 }
723 }
724 }
725
726 if n.min(m).min(p) <= 128 {
727 let transposed: Vec<_> = (0..p)
728 .flat_map(|j| (0..m).map(move |i| b[i * stride_b + j].clone()))
729 .collect();
730 for (a, c) in a.chunks(stride_a).zip(c.chunks_exact_mut(p)) {
731 for (b, c) in transposed.chunks_exact(m).zip(c) {
732 *c = R::dot_product(&a[..m], b);
733 }
734 }
735 return;
736 }
737 let (h, k, w) = (n / 2, m / 2, p / 2);
738 let a11 = 0;
739 let a12 = k;
740 let a21 = h * stride_a;
741 let a22 = a21 + k;
742 let b11 = 0;
743 let b12 = w;
744 let b21 = k * stride_b;
745 let b22 = b21 + w;
746
747 let block = h * w;
748 let mut buf = vec![R::zero(); h * k + k * w + block * 7];
749 let (s1, rest) = buf.split_at_mut(h * k);
750 let (s2, m_buf) = rest.split_at_mut(k * w);
751 let (m1, rest) = m_buf.split_at_mut(block);
752 let (m2, rest) = rest.split_at_mut(block);
753 let (m3, rest) = rest.split_at_mut(block);
754 let (m4, rest) = rest.split_at_mut(block);
755 let (m5, rest) = rest.split_at_mut(block);
756 let (m6, m7) = rest.split_at_mut(block);
757
758 // (A11 + A22)(B11 + B22)
759 add_block::<R>(&a[a11..], &a[a22..], s1, k, stride_a, stride_a);
760 add_block::<R>(&b[b11..], &b[b22..], s2, w, stride_b, stride_b);
761 strassen_rec::<R>(s1, s2, m1, (h, k, w), k, w);
762
763 // (A21 + A22) B11
764 add_block::<R>(&a[a21..], &a[a22..], s1, k, stride_a, stride_a);
765 strassen_rec::<R>(s1, &b[b11..], m2, (h, k, w), k, stride_b);
766
767 // A11 (B12 - B22)
768 sub_block::<R>(&b[b12..], &b[b22..], s2, w, stride_b, stride_b);
769 strassen_rec::<R>(&a[a11..], s2, m3, (h, k, w), stride_a, w);
770
771 // A22 (B21 - B11)
772 sub_block::<R>(&b[b21..], &b[b11..], s2, w, stride_b, stride_b);
773 strassen_rec::<R>(&a[a22..], s2, m4, (h, k, w), stride_a, w);
774
775 // (A11 + A12) B22
776 add_block::<R>(&a[a11..], &a[a12..], s1, k, stride_a, stride_a);
777 strassen_rec::<R>(s1, &b[b22..], m5, (h, k, w), k, stride_b);
778
779 // (A21 - A11)(B11 + B12)
780 sub_block::<R>(&a[a21..], &a[a11..], s1, k, stride_a, stride_a);
781 add_block::<R>(&b[b11..], &b[b12..], s2, w, stride_b, stride_b);
782 strassen_rec::<R>(s1, s2, m6, (h, k, w), k, w);
783
784 // (A12 - A22)(B21 + B22)
785 sub_block::<R>(&a[a12..], &a[a22..], s1, k, stride_a, stride_a);
786 add_block::<R>(&b[b21..], &b[b22..], s2, w, stride_b, stride_b);
787 strassen_rec::<R>(s1, s2, m7, (h, k, w), k, w);
788
789 let c11 = 0;
790 let c12 = w;
791 let c21 = h * p;
792 let c22 = c21 + w;
793 for ((((m1, m4), m5), m7), c) in m1
794 .iter()
795 .zip(m4.iter())
796 .zip(m5.iter())
797 .zip(m7.iter())
798 .zip(c[c11..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
799 {
800 *c = R::add(m1, m4);
801 R::sub_assign(c, m5);
802 R::add_assign(c, m7);
803 }
804 for ((m3, m5), c) in m3
805 .iter()
806 .zip(m5.iter())
807 .zip(c[c12..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
808 {
809 *c = R::add(m3, m5);
810 }
811 for ((m2, m4), c) in m2
812 .iter()
813 .zip(m4.iter())
814 .zip(c[c21..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
815 {
816 *c = R::add(m2, m4);
817 }
818 for ((((m1, m2), m3), m6), c) in m1
819 .iter()
820 .zip(m2.iter())
821 .zip(m3.iter())
822 .zip(m6.iter())
823 .zip(c[c22..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
824 {
825 *c = R::sub(m1, m2);
826 R::add_assign(c, m3);
827 R::add_assign(c, m6);
828 }
829}crates/competitive/src/math/quotient_array.rs (line 103)
84 pub fn min_25_sieve<R>(&self, mut f: impl FnMut(u64, u32) -> T) -> Self
85 where
86 T: Clone + One,
87 R: Ring<T = T, Additive: Invertible>,
88 {
89 let mut dp = self.clone();
90 let max_n = self.isqrtn as u32;
91 with_prime_list(max_n, |pl| {
92 for p in pl.primes_lte(max_n).rev() {
93 let p = u64::from(p);
94 let k = self.quotient_index(p);
95 for (i, q) in Self::index_iter(self.n, self.isqrtn).enumerate() {
96 let mut pc = p;
97 if pc * p > q {
98 break;
99 }
100 let mut c = 1;
101 while q / p >= pc {
102 let x = R::mul(&f(p, c), &(R::sub(&dp[q / pc], &self.data[k])));
103 let x = R::add(&x, &f(p, c + 1));
104 dp.data[i] = R::add(&dp.data[i], &x);
105 c += 1;
106 pc *= p;
107 }
108 }
109 }
110 });
111 for x in &mut dp.data {
112 *x = R::add(x, &T::one());
113 }
114 dp
115 }Sourcefn mul(x: &Self::T, y: &Self::T) -> Self::T
fn mul(x: &Self::T, y: &Self::T) -> Self::T
multiplicative operaion: $+$
Examples found in repository?
More examples
crates/competitive/src/math/bitwiseand_convolve.rs (line 51)
49 fn multiply(f: &mut Self::F, g: &Self::F) {
50 for (f, g) in f.iter_mut().zip(g) {
51 *f = R::mul(f, g);
52 }
53 }
54
55 fn convolve(a: Self::T, b: Self::T) -> Self::T {
56 assert_eq!(a.len(), b.len());
57 let len = a.len();
58 let same = a == b;
59 let mut a = Self::transform(a, len);
60 if same {
61 for a in a.iter_mut() {
62 *a = R::mul(a, a);
63 }
64 } else {
65 let b = Self::transform(b, len);
66 Self::multiply(&mut a, &b);
67 }
68 Self::inverse_transform(a, len)
69 }crates/competitive/src/math/bitwiseor_convolve.rs (line 51)
49 fn multiply(f: &mut Self::F, g: &Self::F) {
50 for (f, g) in f.iter_mut().zip(g) {
51 *f = R::mul(f, g);
52 }
53 }
54
55 fn convolve(a: Self::T, b: Self::T) -> Self::T {
56 assert_eq!(a.len(), b.len());
57 let len = a.len();
58 let same = a == b;
59 let mut a = Self::transform(a, len);
60 if same {
61 for a in a.iter_mut() {
62 *a = R::mul(a, a);
63 }
64 } else {
65 let b = Self::transform(b, len);
66 Self::multiply(&mut a, &b);
67 }
68 Self::inverse_transform(a, len)
69 }crates/competitive/src/algebra/ring.rs (line 48)
45 fn dot_product(x: &[Self::T], y: &[Self::T]) -> Self::T {
46 assert_eq!(x.len(), y.len());
47 x.iter().zip(y).fold(Self::zero(), |mut sum, (x, y)| {
48 Self::add_assign(&mut sum, &Self::mul(x, y));
49 sum
50 })
51 }
52
53 fn add_scaled_assign(x: &mut [Self::T], y: &[Self::T], a: &Self::T) {
54 assert_eq!(x.len(), y.len());
55 for (x, y) in x.iter_mut().zip(y) {
56 Self::add_assign(x, &Self::mul(a, y));
57 }
58 }Additional examples can be found in:
- crates/competitive/src/string/rolling_hash.rs
- crates/competitive/src/algorithm/esper.rs
- crates/competitive/src/math/bitwisexor_convolve.rs
- crates/competitive/src/math/quotient_array.rs
- crates/competitive/src/math/matrix.rs
- crates/competitive/src/math/floor_sum.rs
- crates/competitive/src/math/black_box_matrix.rs
- crates/competitive/src/algorithm/automata_learning.rs
Sourcefn try_matrix_product(
_a: &[Vec<Self::T>],
_b: &[Vec<Self::T>],
) -> Option<Vec<Vec<Self::T>>>
fn try_matrix_product( _a: &[Vec<Self::T>], _b: &[Vec<Self::T>], ) -> Option<Vec<Vec<Self::T>>>
Examples found in repository?
crates/competitive/src/math/matrix.rs (line 669)
667 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
668 assert_eq!(self.shape.1, rhs.shape.0);
669 if let Some(data) = R::try_matrix_product(&self.data, &rhs.data) {
670 return Matrix::from_vec(data);
671 }
672 let rhs = rhs.transpose();
673 Matrix::new_with((self.shape.0, rhs.shape.0), |i, j| {
674 R::dot_product(&self[i], &rhs[j])
675 })
676 }
677}
678
679fn strassen_rec<R: Ring>(
680 a: &[R::T],
681 b: &[R::T],
682 c: &mut [R::T],
683 shape: (usize, usize, usize),
684 stride_a: usize,
685 stride_b: usize,
686) {
687 let (n, m, p) = shape;
688 fn add_block<R: Ring>(
689 a: &[R::T],
690 b: &[R::T],
691 out: &mut [R::T],
692 n: usize,
693 stride_a: usize,
694 stride_b: usize,
695 ) {
696 for ((a, b), c) in a
697 .chunks(stride_a)
698 .zip(b.chunks(stride_b))
699 .zip(out.chunks_exact_mut(n))
700 {
701 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
702 *c = R::add(a, b);
703 }
704 }
705 }
706
707 fn sub_block<R: Ring>(
708 a: &[R::T],
709 b: &[R::T],
710 out: &mut [R::T],
711 n: usize,
712 stride_a: usize,
713 stride_b: usize,
714 ) {
715 for ((a, b), c) in a
716 .chunks(stride_a)
717 .zip(b.chunks(stride_b))
718 .zip(out.chunks_exact_mut(n))
719 {
720 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
721 *c = R::sub(a, b);
722 }
723 }
724 }
725
726 if n.min(m).min(p) <= 128 {
727 let transposed: Vec<_> = (0..p)
728 .flat_map(|j| (0..m).map(move |i| b[i * stride_b + j].clone()))
729 .collect();
730 for (a, c) in a.chunks(stride_a).zip(c.chunks_exact_mut(p)) {
731 for (b, c) in transposed.chunks_exact(m).zip(c) {
732 *c = R::dot_product(&a[..m], b);
733 }
734 }
735 return;
736 }
737 let (h, k, w) = (n / 2, m / 2, p / 2);
738 let a11 = 0;
739 let a12 = k;
740 let a21 = h * stride_a;
741 let a22 = a21 + k;
742 let b11 = 0;
743 let b12 = w;
744 let b21 = k * stride_b;
745 let b22 = b21 + w;
746
747 let block = h * w;
748 let mut buf = vec![R::zero(); h * k + k * w + block * 7];
749 let (s1, rest) = buf.split_at_mut(h * k);
750 let (s2, m_buf) = rest.split_at_mut(k * w);
751 let (m1, rest) = m_buf.split_at_mut(block);
752 let (m2, rest) = rest.split_at_mut(block);
753 let (m3, rest) = rest.split_at_mut(block);
754 let (m4, rest) = rest.split_at_mut(block);
755 let (m5, rest) = rest.split_at_mut(block);
756 let (m6, m7) = rest.split_at_mut(block);
757
758 // (A11 + A22)(B11 + B22)
759 add_block::<R>(&a[a11..], &a[a22..], s1, k, stride_a, stride_a);
760 add_block::<R>(&b[b11..], &b[b22..], s2, w, stride_b, stride_b);
761 strassen_rec::<R>(s1, s2, m1, (h, k, w), k, w);
762
763 // (A21 + A22) B11
764 add_block::<R>(&a[a21..], &a[a22..], s1, k, stride_a, stride_a);
765 strassen_rec::<R>(s1, &b[b11..], m2, (h, k, w), k, stride_b);
766
767 // A11 (B12 - B22)
768 sub_block::<R>(&b[b12..], &b[b22..], s2, w, stride_b, stride_b);
769 strassen_rec::<R>(&a[a11..], s2, m3, (h, k, w), stride_a, w);
770
771 // A22 (B21 - B11)
772 sub_block::<R>(&b[b21..], &b[b11..], s2, w, stride_b, stride_b);
773 strassen_rec::<R>(&a[a22..], s2, m4, (h, k, w), stride_a, w);
774
775 // (A11 + A12) B22
776 add_block::<R>(&a[a11..], &a[a12..], s1, k, stride_a, stride_a);
777 strassen_rec::<R>(s1, &b[b22..], m5, (h, k, w), k, stride_b);
778
779 // (A21 - A11)(B11 + B12)
780 sub_block::<R>(&a[a21..], &a[a11..], s1, k, stride_a, stride_a);
781 add_block::<R>(&b[b11..], &b[b12..], s2, w, stride_b, stride_b);
782 strassen_rec::<R>(s1, s2, m6, (h, k, w), k, w);
783
784 // (A12 - A22)(B21 + B22)
785 sub_block::<R>(&a[a12..], &a[a22..], s1, k, stride_a, stride_a);
786 add_block::<R>(&b[b21..], &b[b22..], s2, w, stride_b, stride_b);
787 strassen_rec::<R>(s1, s2, m7, (h, k, w), k, w);
788
789 let c11 = 0;
790 let c12 = w;
791 let c21 = h * p;
792 let c22 = c21 + w;
793 for ((((m1, m4), m5), m7), c) in m1
794 .iter()
795 .zip(m4.iter())
796 .zip(m5.iter())
797 .zip(m7.iter())
798 .zip(c[c11..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
799 {
800 *c = R::add(m1, m4);
801 R::sub_assign(c, m5);
802 R::add_assign(c, m7);
803 }
804 for ((m3, m5), c) in m3
805 .iter()
806 .zip(m5.iter())
807 .zip(c[c12..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
808 {
809 *c = R::add(m3, m5);
810 }
811 for ((m2, m4), c) in m2
812 .iter()
813 .zip(m4.iter())
814 .zip(c[c21..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
815 {
816 *c = R::add(m2, m4);
817 }
818 for ((((m1, m2), m3), m6), c) in m1
819 .iter()
820 .zip(m2.iter())
821 .zip(m3.iter())
822 .zip(m6.iter())
823 .zip(c[c22..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
824 {
825 *c = R::sub(m1, m2);
826 R::add_assign(c, m3);
827 R::add_assign(c, m6);
828 }
829}
830
831impl<R> Matrix<R>
832where
833 R: Ring,
834{
835 pub fn mul_strassen(&self, rhs: &Matrix<R>) -> Matrix<R> {
836 assert_eq!(self.shape.1, rhs.shape.0);
837 if let Some(data) = R::try_matrix_product(&self.data, &rhs.data) {
838 return Matrix::from_vec(data);
839 }
840 let (n, m) = self.shape;
841 let p = rhs.shape.1;
842 if n == 0 || m == 0 || p == 0 {
843 return Matrix::zeros((n, p));
844 }
845 let split = n.min(m).min(p).div_ceil(128).next_power_of_two();
846 if split <= 2 {
847 return self * rhs;
848 }
849 let rows = n.div_ceil(split) * split;
850 let inner = m.div_ceil(split) * split;
851 let cols = p.div_ceil(split) * split;
852 let mut a = vec![R::zero(); rows * inner];
853 for (a, data) in a.chunks_exact_mut(inner).zip(&self.data) {
854 a[..m].clone_from_slice(data);
855 }
856 let mut b = vec![R::zero(); inner * cols];
857 for (b, data) in b.chunks_exact_mut(cols).zip(&rhs.data) {
858 b[..p].clone_from_slice(data);
859 }
860 let mut c = vec![R::zero(); rows * cols];
861 strassen_rec::<R>(&a, &b, &mut c, (rows, inner, cols), inner, cols);
862 let mut res = Matrix::zeros((n, p));
863 for (data, c) in res.data.iter_mut().zip(c.chunks_exact(cols)) {
864 data.clone_from_slice(&c[..p]);
865 }
866 res
867 }Sourcefn dot_product(x: &[Self::T], y: &[Self::T]) -> Self::T
fn dot_product(x: &[Self::T], y: &[Self::T]) -> Self::T
Examples found in repository?
More examples
crates/competitive/src/math/subset_convolve.rs (line 45)
32 fn multiply_row(
33 x: &[R::T],
34 y: &[R::T],
35 right: &mut [R::T],
36 output: &mut [R::T],
37 rank: usize,
38 ) -> usize {
39 for (right, y) in right[..=rank].iter_mut().zip(y[..=rank].iter().rev()) {
40 right.clone_from(y);
41 }
42 let end = (rank * 2).min(x.len() - 1);
43 for (degree, output) in output.iter_mut().enumerate().take(end + 1).skip(rank) {
44 let first = degree - rank;
45 *output = R::dot_product(&x[first..=rank], &right[..=rank - first]);
46 }
47 end
48 }crates/competitive/src/math/matrix.rs (line 178)
158 fn eliminate<const DETERMINANT: bool>(&mut self) -> (usize, R::T) {
159 let (n, m) = self.shape;
160 let mut rank = 0;
161 let mut determinant = R::one();
162 let mut negative = false;
163 for first in (0..m).step_by(64) {
164 if rank == n {
165 break;
166 }
167 let end = (first + 64).min(m);
168 let start = rank;
169 let mut pivots = Vec::new();
170 let mut panel = vec![vec![R::zero(); end - first]; end - first];
171 for col in first..end {
172 if panel[col - first][..col - first]
173 .iter()
174 .any(|x| !R::is_zero(x))
175 {
176 for row in &mut self.data[rank..] {
177 let value =
178 R::dot_product(&row[first..col], &panel[col - first][..col - first]);
179 R::sub_assign(&mut row[col], &value);
180 }
181 }
182 let Some(pivot) = (rank..n).find(|&i| !R::is_zero(&self[i][col])) else {
183 continue;
184 };
185 if pivot != rank {
186 self.data.swap(rank, pivot);
187 negative = !negative;
188 }
189 if DETERMINANT {
190 R::mul_assign(&mut determinant, &self[rank][col]);
191 }
192 let inv = R::inv(&self[rank][col]);
193 let row = &mut self.data[rank];
194 for c in col + 1..end {
195 let value = R::dot_product(&row[first..col], &panel[c - first][..col - first]);
196 R::sub_assign(&mut row[c], &value);
197 panel[c - first][col - first] = row[c].clone();
198 }
199 for row in &mut self.data[rank + 1..] {
200 R::mul_assign(&mut row[col], &inv);
201 }
202 pivots.push(col);
203 rank += 1;
204 if rank == n {
205 break;
206 }
207 }
208 for i in start..if rank - start < 32 { n } else { rank } {
209 let (upper, lower) = self.data.split_at_mut(i);
210 let row = &mut lower[0];
211 for (j, &col) in pivots[..(i - start).min(pivots.len())].iter().enumerate() {
212 if R::is_zero(&row[col]) {
213 continue;
214 }
215 let factor = R::neg(&row[col]);
216 R::add_scaled_assign(&mut row[end..], &upper[start + j][end..], &factor);
217 }
218 }
219 if rank < n
220 && end < m
221 && rank - start >= 32
222 && self.data[rank..]
223 .iter()
224 .any(|row| pivots.iter().any(|&col| !R::is_zero(&row[col])))
225 {
226 let lower = Self::new_with((n - rank, rank - start), |i, j| {
227 R::neg(&self[rank + i][pivots[j]])
228 });
229 let upper = Self::new_with((rank - start, m - end), |i, j| {
230 self[start + i][end + j].clone()
231 });
232 let update = &lower * &upper;
233 for (row, update) in self.data[rank..].iter_mut().zip(update.data) {
234 for (x, y) in row[end..].iter_mut().zip(update) {
235 R::add_assign(x, &y);
236 }
237 }
238 }
239 for (i, &col) in pivots.iter().enumerate() {
240 for row in &mut self.data[start + i + 1..] {
241 row[col] = R::zero();
242 }
243 }
244 if DETERMINANT && rank < end {
245 return (rank, R::zero());
246 }
247 }
248 if DETERMINANT && negative {
249 determinant = R::neg(&determinant);
250 }
251 (rank, determinant)
252 }
253
254 /// f: (row, pivot_row, col)
255 pub fn row_reduction_with<F>(&mut self, normalize: bool, mut f: F)
256 where
257 F: FnMut(usize, usize, usize),
258 {
259 let (n, m) = self.shape;
260 let mut c = 0;
261 for r in 0..n {
262 loop {
263 if c >= m {
264 return;
265 }
266 if let Some(pivot) = (r..n).find(|&p| !R::is_zero(&self[p][c])) {
267 f(r, pivot, c);
268 self.data.swap(r, pivot);
269 break;
270 };
271 c += 1;
272 }
273 let d = R::inv(&self[r][c]);
274 if normalize {
275 for value in &mut self[r][c..m] {
276 R::mul_assign(value, &d);
277 }
278 }
279 for i in (0..n).filter(|&i| i != r) {
280 let mut e = self[i][c].clone();
281 if !normalize {
282 R::mul_assign(&mut e, &d);
283 }
284 for j in c..m {
285 let e = R::mul(&e, &self[r][j]);
286 R::sub_assign(&mut self[i][j], &e);
287 }
288 }
289 c += 1;
290 }
291 }
292
293 pub fn row_reduction(&mut self, normalize: bool) {
294 self.row_reduction_with(normalize, |_, _, _| {});
295 }
296
297 pub fn rank(&mut self) -> usize {
298 self.eliminate::<false>().0
299 }
300
301 pub fn determinant(&mut self) -> R::T {
302 assert_eq!(self.shape.0, self.shape.1);
303 self.eliminate::<true>().1
304 }
305
306 pub fn solve_system_of_linear_equations(
307 &self,
308 b: &[R::T],
309 ) -> Option<SystemOfLinearEquationsSolution<R>> {
310 assert_eq!(self.shape.0, b.len());
311 let m = self.shape.1;
312 let mut a = Self::new_with((self.shape.0, m + 1), |i, j| {
313 if j == m {
314 b[i].clone()
315 } else {
316 self[i][j].clone()
317 }
318 });
319 let rank = a.eliminate::<false>().0;
320 let mut pivots = Vec::with_capacity(rank);
321 let mut b = Vec::with_capacity(rank);
322 for row in &a.data[..rank] {
323 let c = row.iter().position(|x| !R::is_zero(x)).unwrap();
324 if c == m {
325 return None;
326 }
327 pivots.push(c);
328 b.push(row[m].clone());
329 }
330
331 let mut free = Vec::with_capacity(m - rank);
332 let mut pivot = 0;
333 for c in 0..m {
334 if pivot < rank && pivots[pivot] == c {
335 pivot += 1;
336 } else {
337 free.push(c);
338 }
339 }
340 let mut coefficients: Vec<Vec<_>> = (0..rank)
341 .map(|i| free.iter().map(|&c| a[i][c].clone()).collect())
342 .collect();
343 for k in (0..rank).rev() {
344 let c = pivots[k];
345 let inv = R::inv(&a[k][c]);
346 R::mul_assign(&mut b[k], &inv);
347 let pivot_b = b[k].clone();
348 let (upper, lower) = coefficients.split_at_mut(k);
349 let pivot_coefficients = &mut lower[0];
350 for x in pivot_coefficients.iter_mut() {
351 R::mul_assign(x, &inv);
352 }
353 for ((row, value), coefficients) in a.data[..k].iter_mut().zip(&mut b[..k]).zip(upper) {
354 if R::is_zero(&row[c]) {
355 continue;
356 }
357 let factor = row[c].clone();
358 row[c] = R::zero();
359 R::sub_assign(value, &R::mul(&factor, &pivot_b));
360 R::add_scaled_assign(coefficients, pivot_coefficients, &R::neg(&factor));
361 }
362 }
363
364 let mut particular = vec![R::zero(); m];
365 for i in 0..rank {
366 particular[pivots[i]] = b[i].clone();
367 }
368 let mut basis = Vec::with_capacity(free.len());
369 for (j, &c) in free.iter().enumerate() {
370 let mut vector = vec![R::zero(); m];
371 vector[c] = R::one();
372 for i in 0..rank {
373 vector[pivots[i]] = R::neg(&coefficients[i][j]);
374 }
375 basis.push(vector);
376 }
377 Some(SystemOfLinearEquationsSolution { particular, basis })
378 }
379
380 pub fn inverse(&self) -> Option<Matrix<R>> {
381 assert_eq!(self.shape.0, self.shape.1);
382 let n = self.shape.0;
383 if n >= 64 {
384 let m = n / 2;
385 let a = Self::new_with((m, m), |i, j| self[i][j].clone());
386 if let Some(mut ai) = a.inverse() {
387 let b = Self::new_with((m, n - m), |i, j| self[i][j + m].clone());
388 let c = Self::new_with((n - m, m), |i, j| self[i + m][j].clone());
389 let mut d = Self::new_with((n - m, n - m), |i, j| self[i + m][j + m].clone());
390 let u = &ai * &b;
391 let v = &c * &ai;
392 d -= &v * &b;
393 let di = d.inverse()?;
394 let r = &u * &di;
395 let t = &di * &v;
396 ai += &r * &v;
397 let mut inverse = Self::zeros((n, n));
398 for i in 0..m {
399 inverse[i][..m].clone_from_slice(&ai[i]);
400 for (x, y) in inverse[i][m..].iter_mut().zip(&r[i]) {
401 *x = R::neg(y);
402 }
403 }
404 for i in m..n {
405 for (x, y) in inverse[i][..m].iter_mut().zip(&t[i - m]) {
406 *x = R::neg(y);
407 }
408 inverse[i][m..].clone_from_slice(&di[i - m]);
409 }
410 return Some(inverse);
411 }
412 }
413 let mut a = self.clone();
414 let mut inverse = Self::eye((n, n));
415 let mut ranges: Vec<_> = (0..n).map(|i| (i, i + 1)).collect();
416 for r in 0..n {
417 let pivot = (r..n).find(|&i| !R::is_zero(&a[i][r]))?;
418 a.data.swap(r, pivot);
419 inverse.data.swap(r, pivot);
420 ranges.swap(r, pivot);
421
422 let d = R::inv(&a[r][r]);
423 for x in &mut a[r][r..] {
424 R::mul_assign(x, &d);
425 }
426 let (left, right) = ranges[r];
427 for x in &mut inverse[r][left..right] {
428 R::mul_assign(x, &d);
429 }
430
431 let (a_upper, a_lower) = a.data.split_at_mut(r + 1);
432 let pivot_a = &a_upper[r];
433 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r + 1);
434 let pivot_inverse = &inverse_upper[r];
435 let (ranges_upper, ranges_lower) = ranges.split_at_mut(r + 1);
436 let (left, right) = ranges_upper[r];
437 for ((a, inverse), range) in a_lower.iter_mut().zip(inverse_lower).zip(ranges_lower) {
438 if R::is_zero(&a[r]) {
439 continue;
440 }
441 let e = a[r].clone();
442 a[r] = R::zero();
443 R::add_scaled_assign(&mut a[(r + 1)..], &pivot_a[(r + 1)..], &R::neg(&e));
444 R::add_scaled_assign(
445 &mut inverse[left..right],
446 &pivot_inverse[left..right],
447 &R::neg(&e),
448 );
449 range.0 = range.0.min(left);
450 range.1 = range.1.max(right);
451 }
452 }
453 for r in (0..n).rev() {
454 let (left, right) = ranges[r];
455 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r);
456 let pivot_inverse = &inverse_lower[0];
457 let (ranges_upper, _) = ranges.split_at_mut(r);
458 for ((a, inverse), range) in a.data[..r].iter_mut().zip(inverse_upper).zip(ranges_upper)
459 {
460 if R::is_zero(&a[r]) {
461 continue;
462 }
463 let e = a[r].clone();
464 a[r] = R::zero();
465 R::add_scaled_assign(
466 &mut inverse[left..right],
467 &pivot_inverse[left..right],
468 &R::neg(&e),
469 );
470 range.0 = range.0.min(left);
471 range.1 = range.1.max(right);
472 }
473 }
474 Some(inverse)
475 }
476
477 pub fn characteristic_polynomial(&mut self) -> Vec<R::T> {
478 let n = self.shape.0;
479 if n == 0 {
480 return vec![R::one()];
481 }
482 assert!(self.data.iter().all(|a| a.len() == n));
483 for j in 0..(n - 1) {
484 if let Some(x) = ((j + 1)..n).find(|&x| !R::is_zero(&self[x][j])) {
485 self.data.swap(j + 1, x);
486 self.data.iter_mut().for_each(|a| a.swap(j + 1, x));
487 let inv = R::inv(&self[j + 1][j]);
488 let mut v = vec![];
489 let src = std::mem::take(&mut self[j + 1]);
490 for a in self.data[(j + 2)..].iter_mut() {
491 let mul = R::mul(&a[j], &inv);
492 R::add_scaled_assign(&mut a[j..], &src[j..], &R::neg(&mul));
493 v.push(mul);
494 }
495 self[j + 1] = src;
496 for a in self.data.iter_mut() {
497 let v = R::dot_product(&a[(j + 2)..], &v);
498 R::add_assign(&mut a[j + 1], &v);
499 }
500 }
501 }
502 // dp[k][j - k] stores [x^k] det(xI - A[..j, ..j]).
503 let mut dp: Vec<Vec<R::T>> = (0..=n).map(|i| Vec::with_capacity(n + 1 - i)).collect();
504 dp[0].push(R::one());
505 for i in 0..n {
506 let mut c = vec![R::zero(); i + 1];
507 c[i] = R::neg(&self[i][i]);
508 let mut mul = R::one();
509 for j in (0..i).rev() {
510 mul = R::mul(&mul, &self[j + 1][j]);
511 c[j] = R::neg(&R::mul(&mul, &self[j][i]));
512 }
513 for k in (0..=i).rev() {
514 let mut value = R::dot_product(&dp[k], &c[k..]);
515 if k > 0 {
516 R::add_assign(&mut value, dp[k - 1].last().unwrap());
517 }
518 dp[k].push(value);
519 }
520 dp[i + 1].push(R::one());
521 }
522 dp.into_iter().map(|mut c| c.pop().unwrap()).collect()
523 }
524}
525
526impl<R> Index<usize> for Matrix<R>
527where
528 R: SemiRing,
529{
530 type Output = Vec<R::T>;
531 fn index(&self, index: usize) -> &Self::Output {
532 &self.data[index]
533 }
534}
535
536impl<R> IndexMut<usize> for Matrix<R>
537where
538 R: SemiRing,
539{
540 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
541 &mut self.data[index]
542 }
543}
544
545impl<R> Index<(usize, usize)> for Matrix<R>
546where
547 R: SemiRing,
548{
549 type Output = R::T;
550 fn index(&self, index: (usize, usize)) -> &Self::Output {
551 &self.data[index.0][index.1]
552 }
553}
554
555impl<R> IndexMut<(usize, usize)> for Matrix<R>
556where
557 R: SemiRing,
558{
559 fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
560 &mut self.data[index.0][index.1]
561 }
562}
563
564macro_rules! impl_matrix_pairwise_binop {
565 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident $(where [$($clauses:tt)*])?) => {
566 impl<R> $imp_assign for Matrix<R>
567 where
568 R: SemiRing,
569 $($($clauses)*)?
570 {
571 fn $method_assign(&mut self, rhs: Self) {
572 self.pairwise_assign(&rhs, |a, b| R::$method_assign(a, b));
573 }
574 }
575 impl<R> $imp_assign<&Matrix<R>> for Matrix<R>
576 where
577 R: SemiRing,
578 $($($clauses)*)?
579 {
580 fn $method_assign(&mut self, rhs: &Self) {
581 self.pairwise_assign(rhs, |a, b| R::$method_assign(a, b));
582 }
583 }
584 impl<R> $imp for Matrix<R>
585 where
586 R: SemiRing,
587 $($($clauses)*)?
588 {
589 type Output = Matrix<R>;
590 fn $method(mut self, rhs: Self) -> Self::Output {
591 self.$method_assign(rhs);
592 self
593 }
594 }
595 impl<R> $imp<&Matrix<R>> for Matrix<R>
596 where
597 R: SemiRing,
598 $($($clauses)*)?
599 {
600 type Output = Matrix<R>;
601 fn $method(mut self, rhs: &Self) -> Self::Output {
602 self.$method_assign(rhs);
603 self
604 }
605 }
606 impl<R> $imp<Matrix<R>> for &Matrix<R>
607 where
608 R: SemiRing,
609 $($($clauses)*)?
610 {
611 type Output = Matrix<R>;
612 fn $method(self, mut rhs: Matrix<R>) -> Self::Output {
613 rhs.pairwise_assign(self, |a, b| *a = R::$method(b, a));
614 rhs
615 }
616 }
617 impl<R> $imp<&Matrix<R>> for &Matrix<R>
618 where
619 R: SemiRing,
620 $($($clauses)*)?
621 {
622 type Output = Matrix<R>;
623 fn $method(self, rhs: &Matrix<R>) -> Self::Output {
624 let mut this = self.clone();
625 this.$method_assign(rhs);
626 this
627 }
628 }
629 };
630}
631
632impl_matrix_pairwise_binop!(Add, add, AddAssign, add_assign);
633impl_matrix_pairwise_binop!(Sub, sub, SubAssign, sub_assign where [R: SemiRing<Additive: Invertible>]);
634
635impl<R> Mul for Matrix<R>
636where
637 R: SemiRing,
638{
639 type Output = Matrix<R>;
640 fn mul(self, rhs: Self) -> Self::Output {
641 (&self).mul(&rhs)
642 }
643}
644impl<R> Mul<&Matrix<R>> for Matrix<R>
645where
646 R: SemiRing,
647{
648 type Output = Matrix<R>;
649 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
650 (&self).mul(rhs)
651 }
652}
653impl<R> Mul<Matrix<R>> for &Matrix<R>
654where
655 R: SemiRing,
656{
657 type Output = Matrix<R>;
658 fn mul(self, rhs: Matrix<R>) -> Self::Output {
659 self.mul(&rhs)
660 }
661}
662impl<R> Mul<&Matrix<R>> for &Matrix<R>
663where
664 R: SemiRing,
665{
666 type Output = Matrix<R>;
667 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
668 assert_eq!(self.shape.1, rhs.shape.0);
669 if let Some(data) = R::try_matrix_product(&self.data, &rhs.data) {
670 return Matrix::from_vec(data);
671 }
672 let rhs = rhs.transpose();
673 Matrix::new_with((self.shape.0, rhs.shape.0), |i, j| {
674 R::dot_product(&self[i], &rhs[j])
675 })
676 }
677}
678
679fn strassen_rec<R: Ring>(
680 a: &[R::T],
681 b: &[R::T],
682 c: &mut [R::T],
683 shape: (usize, usize, usize),
684 stride_a: usize,
685 stride_b: usize,
686) {
687 let (n, m, p) = shape;
688 fn add_block<R: Ring>(
689 a: &[R::T],
690 b: &[R::T],
691 out: &mut [R::T],
692 n: usize,
693 stride_a: usize,
694 stride_b: usize,
695 ) {
696 for ((a, b), c) in a
697 .chunks(stride_a)
698 .zip(b.chunks(stride_b))
699 .zip(out.chunks_exact_mut(n))
700 {
701 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
702 *c = R::add(a, b);
703 }
704 }
705 }
706
707 fn sub_block<R: Ring>(
708 a: &[R::T],
709 b: &[R::T],
710 out: &mut [R::T],
711 n: usize,
712 stride_a: usize,
713 stride_b: usize,
714 ) {
715 for ((a, b), c) in a
716 .chunks(stride_a)
717 .zip(b.chunks(stride_b))
718 .zip(out.chunks_exact_mut(n))
719 {
720 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
721 *c = R::sub(a, b);
722 }
723 }
724 }
725
726 if n.min(m).min(p) <= 128 {
727 let transposed: Vec<_> = (0..p)
728 .flat_map(|j| (0..m).map(move |i| b[i * stride_b + j].clone()))
729 .collect();
730 for (a, c) in a.chunks(stride_a).zip(c.chunks_exact_mut(p)) {
731 for (b, c) in transposed.chunks_exact(m).zip(c) {
732 *c = R::dot_product(&a[..m], b);
733 }
734 }
735 return;
736 }
737 let (h, k, w) = (n / 2, m / 2, p / 2);
738 let a11 = 0;
739 let a12 = k;
740 let a21 = h * stride_a;
741 let a22 = a21 + k;
742 let b11 = 0;
743 let b12 = w;
744 let b21 = k * stride_b;
745 let b22 = b21 + w;
746
747 let block = h * w;
748 let mut buf = vec![R::zero(); h * k + k * w + block * 7];
749 let (s1, rest) = buf.split_at_mut(h * k);
750 let (s2, m_buf) = rest.split_at_mut(k * w);
751 let (m1, rest) = m_buf.split_at_mut(block);
752 let (m2, rest) = rest.split_at_mut(block);
753 let (m3, rest) = rest.split_at_mut(block);
754 let (m4, rest) = rest.split_at_mut(block);
755 let (m5, rest) = rest.split_at_mut(block);
756 let (m6, m7) = rest.split_at_mut(block);
757
758 // (A11 + A22)(B11 + B22)
759 add_block::<R>(&a[a11..], &a[a22..], s1, k, stride_a, stride_a);
760 add_block::<R>(&b[b11..], &b[b22..], s2, w, stride_b, stride_b);
761 strassen_rec::<R>(s1, s2, m1, (h, k, w), k, w);
762
763 // (A21 + A22) B11
764 add_block::<R>(&a[a21..], &a[a22..], s1, k, stride_a, stride_a);
765 strassen_rec::<R>(s1, &b[b11..], m2, (h, k, w), k, stride_b);
766
767 // A11 (B12 - B22)
768 sub_block::<R>(&b[b12..], &b[b22..], s2, w, stride_b, stride_b);
769 strassen_rec::<R>(&a[a11..], s2, m3, (h, k, w), stride_a, w);
770
771 // A22 (B21 - B11)
772 sub_block::<R>(&b[b21..], &b[b11..], s2, w, stride_b, stride_b);
773 strassen_rec::<R>(&a[a22..], s2, m4, (h, k, w), stride_a, w);
774
775 // (A11 + A12) B22
776 add_block::<R>(&a[a11..], &a[a12..], s1, k, stride_a, stride_a);
777 strassen_rec::<R>(s1, &b[b22..], m5, (h, k, w), k, stride_b);
778
779 // (A21 - A11)(B11 + B12)
780 sub_block::<R>(&a[a21..], &a[a11..], s1, k, stride_a, stride_a);
781 add_block::<R>(&b[b11..], &b[b12..], s2, w, stride_b, stride_b);
782 strassen_rec::<R>(s1, s2, m6, (h, k, w), k, w);
783
784 // (A12 - A22)(B21 + B22)
785 sub_block::<R>(&a[a12..], &a[a22..], s1, k, stride_a, stride_a);
786 add_block::<R>(&b[b21..], &b[b22..], s2, w, stride_b, stride_b);
787 strassen_rec::<R>(s1, s2, m7, (h, k, w), k, w);
788
789 let c11 = 0;
790 let c12 = w;
791 let c21 = h * p;
792 let c22 = c21 + w;
793 for ((((m1, m4), m5), m7), c) in m1
794 .iter()
795 .zip(m4.iter())
796 .zip(m5.iter())
797 .zip(m7.iter())
798 .zip(c[c11..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
799 {
800 *c = R::add(m1, m4);
801 R::sub_assign(c, m5);
802 R::add_assign(c, m7);
803 }
804 for ((m3, m5), c) in m3
805 .iter()
806 .zip(m5.iter())
807 .zip(c[c12..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
808 {
809 *c = R::add(m3, m5);
810 }
811 for ((m2, m4), c) in m2
812 .iter()
813 .zip(m4.iter())
814 .zip(c[c21..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
815 {
816 *c = R::add(m2, m4);
817 }
818 for ((((m1, m2), m3), m6), c) in m1
819 .iter()
820 .zip(m2.iter())
821 .zip(m3.iter())
822 .zip(m6.iter())
823 .zip(c[c22..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
824 {
825 *c = R::sub(m1, m2);
826 R::add_assign(c, m3);
827 R::add_assign(c, m6);
828 }
829}Sourcefn add_scaled_assign(x: &mut [Self::T], y: &[Self::T], a: &Self::T)
fn add_scaled_assign(x: &mut [Self::T], y: &[Self::T], a: &Self::T)
Examples found in repository?
crates/competitive/src/math/matrix.rs (line 216)
158 fn eliminate<const DETERMINANT: bool>(&mut self) -> (usize, R::T) {
159 let (n, m) = self.shape;
160 let mut rank = 0;
161 let mut determinant = R::one();
162 let mut negative = false;
163 for first in (0..m).step_by(64) {
164 if rank == n {
165 break;
166 }
167 let end = (first + 64).min(m);
168 let start = rank;
169 let mut pivots = Vec::new();
170 let mut panel = vec![vec![R::zero(); end - first]; end - first];
171 for col in first..end {
172 if panel[col - first][..col - first]
173 .iter()
174 .any(|x| !R::is_zero(x))
175 {
176 for row in &mut self.data[rank..] {
177 let value =
178 R::dot_product(&row[first..col], &panel[col - first][..col - first]);
179 R::sub_assign(&mut row[col], &value);
180 }
181 }
182 let Some(pivot) = (rank..n).find(|&i| !R::is_zero(&self[i][col])) else {
183 continue;
184 };
185 if pivot != rank {
186 self.data.swap(rank, pivot);
187 negative = !negative;
188 }
189 if DETERMINANT {
190 R::mul_assign(&mut determinant, &self[rank][col]);
191 }
192 let inv = R::inv(&self[rank][col]);
193 let row = &mut self.data[rank];
194 for c in col + 1..end {
195 let value = R::dot_product(&row[first..col], &panel[c - first][..col - first]);
196 R::sub_assign(&mut row[c], &value);
197 panel[c - first][col - first] = row[c].clone();
198 }
199 for row in &mut self.data[rank + 1..] {
200 R::mul_assign(&mut row[col], &inv);
201 }
202 pivots.push(col);
203 rank += 1;
204 if rank == n {
205 break;
206 }
207 }
208 for i in start..if rank - start < 32 { n } else { rank } {
209 let (upper, lower) = self.data.split_at_mut(i);
210 let row = &mut lower[0];
211 for (j, &col) in pivots[..(i - start).min(pivots.len())].iter().enumerate() {
212 if R::is_zero(&row[col]) {
213 continue;
214 }
215 let factor = R::neg(&row[col]);
216 R::add_scaled_assign(&mut row[end..], &upper[start + j][end..], &factor);
217 }
218 }
219 if rank < n
220 && end < m
221 && rank - start >= 32
222 && self.data[rank..]
223 .iter()
224 .any(|row| pivots.iter().any(|&col| !R::is_zero(&row[col])))
225 {
226 let lower = Self::new_with((n - rank, rank - start), |i, j| {
227 R::neg(&self[rank + i][pivots[j]])
228 });
229 let upper = Self::new_with((rank - start, m - end), |i, j| {
230 self[start + i][end + j].clone()
231 });
232 let update = &lower * &upper;
233 for (row, update) in self.data[rank..].iter_mut().zip(update.data) {
234 for (x, y) in row[end..].iter_mut().zip(update) {
235 R::add_assign(x, &y);
236 }
237 }
238 }
239 for (i, &col) in pivots.iter().enumerate() {
240 for row in &mut self.data[start + i + 1..] {
241 row[col] = R::zero();
242 }
243 }
244 if DETERMINANT && rank < end {
245 return (rank, R::zero());
246 }
247 }
248 if DETERMINANT && negative {
249 determinant = R::neg(&determinant);
250 }
251 (rank, determinant)
252 }
253
254 /// f: (row, pivot_row, col)
255 pub fn row_reduction_with<F>(&mut self, normalize: bool, mut f: F)
256 where
257 F: FnMut(usize, usize, usize),
258 {
259 let (n, m) = self.shape;
260 let mut c = 0;
261 for r in 0..n {
262 loop {
263 if c >= m {
264 return;
265 }
266 if let Some(pivot) = (r..n).find(|&p| !R::is_zero(&self[p][c])) {
267 f(r, pivot, c);
268 self.data.swap(r, pivot);
269 break;
270 };
271 c += 1;
272 }
273 let d = R::inv(&self[r][c]);
274 if normalize {
275 for value in &mut self[r][c..m] {
276 R::mul_assign(value, &d);
277 }
278 }
279 for i in (0..n).filter(|&i| i != r) {
280 let mut e = self[i][c].clone();
281 if !normalize {
282 R::mul_assign(&mut e, &d);
283 }
284 for j in c..m {
285 let e = R::mul(&e, &self[r][j]);
286 R::sub_assign(&mut self[i][j], &e);
287 }
288 }
289 c += 1;
290 }
291 }
292
293 pub fn row_reduction(&mut self, normalize: bool) {
294 self.row_reduction_with(normalize, |_, _, _| {});
295 }
296
297 pub fn rank(&mut self) -> usize {
298 self.eliminate::<false>().0
299 }
300
301 pub fn determinant(&mut self) -> R::T {
302 assert_eq!(self.shape.0, self.shape.1);
303 self.eliminate::<true>().1
304 }
305
306 pub fn solve_system_of_linear_equations(
307 &self,
308 b: &[R::T],
309 ) -> Option<SystemOfLinearEquationsSolution<R>> {
310 assert_eq!(self.shape.0, b.len());
311 let m = self.shape.1;
312 let mut a = Self::new_with((self.shape.0, m + 1), |i, j| {
313 if j == m {
314 b[i].clone()
315 } else {
316 self[i][j].clone()
317 }
318 });
319 let rank = a.eliminate::<false>().0;
320 let mut pivots = Vec::with_capacity(rank);
321 let mut b = Vec::with_capacity(rank);
322 for row in &a.data[..rank] {
323 let c = row.iter().position(|x| !R::is_zero(x)).unwrap();
324 if c == m {
325 return None;
326 }
327 pivots.push(c);
328 b.push(row[m].clone());
329 }
330
331 let mut free = Vec::with_capacity(m - rank);
332 let mut pivot = 0;
333 for c in 0..m {
334 if pivot < rank && pivots[pivot] == c {
335 pivot += 1;
336 } else {
337 free.push(c);
338 }
339 }
340 let mut coefficients: Vec<Vec<_>> = (0..rank)
341 .map(|i| free.iter().map(|&c| a[i][c].clone()).collect())
342 .collect();
343 for k in (0..rank).rev() {
344 let c = pivots[k];
345 let inv = R::inv(&a[k][c]);
346 R::mul_assign(&mut b[k], &inv);
347 let pivot_b = b[k].clone();
348 let (upper, lower) = coefficients.split_at_mut(k);
349 let pivot_coefficients = &mut lower[0];
350 for x in pivot_coefficients.iter_mut() {
351 R::mul_assign(x, &inv);
352 }
353 for ((row, value), coefficients) in a.data[..k].iter_mut().zip(&mut b[..k]).zip(upper) {
354 if R::is_zero(&row[c]) {
355 continue;
356 }
357 let factor = row[c].clone();
358 row[c] = R::zero();
359 R::sub_assign(value, &R::mul(&factor, &pivot_b));
360 R::add_scaled_assign(coefficients, pivot_coefficients, &R::neg(&factor));
361 }
362 }
363
364 let mut particular = vec![R::zero(); m];
365 for i in 0..rank {
366 particular[pivots[i]] = b[i].clone();
367 }
368 let mut basis = Vec::with_capacity(free.len());
369 for (j, &c) in free.iter().enumerate() {
370 let mut vector = vec![R::zero(); m];
371 vector[c] = R::one();
372 for i in 0..rank {
373 vector[pivots[i]] = R::neg(&coefficients[i][j]);
374 }
375 basis.push(vector);
376 }
377 Some(SystemOfLinearEquationsSolution { particular, basis })
378 }
379
380 pub fn inverse(&self) -> Option<Matrix<R>> {
381 assert_eq!(self.shape.0, self.shape.1);
382 let n = self.shape.0;
383 if n >= 64 {
384 let m = n / 2;
385 let a = Self::new_with((m, m), |i, j| self[i][j].clone());
386 if let Some(mut ai) = a.inverse() {
387 let b = Self::new_with((m, n - m), |i, j| self[i][j + m].clone());
388 let c = Self::new_with((n - m, m), |i, j| self[i + m][j].clone());
389 let mut d = Self::new_with((n - m, n - m), |i, j| self[i + m][j + m].clone());
390 let u = &ai * &b;
391 let v = &c * &ai;
392 d -= &v * &b;
393 let di = d.inverse()?;
394 let r = &u * &di;
395 let t = &di * &v;
396 ai += &r * &v;
397 let mut inverse = Self::zeros((n, n));
398 for i in 0..m {
399 inverse[i][..m].clone_from_slice(&ai[i]);
400 for (x, y) in inverse[i][m..].iter_mut().zip(&r[i]) {
401 *x = R::neg(y);
402 }
403 }
404 for i in m..n {
405 for (x, y) in inverse[i][..m].iter_mut().zip(&t[i - m]) {
406 *x = R::neg(y);
407 }
408 inverse[i][m..].clone_from_slice(&di[i - m]);
409 }
410 return Some(inverse);
411 }
412 }
413 let mut a = self.clone();
414 let mut inverse = Self::eye((n, n));
415 let mut ranges: Vec<_> = (0..n).map(|i| (i, i + 1)).collect();
416 for r in 0..n {
417 let pivot = (r..n).find(|&i| !R::is_zero(&a[i][r]))?;
418 a.data.swap(r, pivot);
419 inverse.data.swap(r, pivot);
420 ranges.swap(r, pivot);
421
422 let d = R::inv(&a[r][r]);
423 for x in &mut a[r][r..] {
424 R::mul_assign(x, &d);
425 }
426 let (left, right) = ranges[r];
427 for x in &mut inverse[r][left..right] {
428 R::mul_assign(x, &d);
429 }
430
431 let (a_upper, a_lower) = a.data.split_at_mut(r + 1);
432 let pivot_a = &a_upper[r];
433 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r + 1);
434 let pivot_inverse = &inverse_upper[r];
435 let (ranges_upper, ranges_lower) = ranges.split_at_mut(r + 1);
436 let (left, right) = ranges_upper[r];
437 for ((a, inverse), range) in a_lower.iter_mut().zip(inverse_lower).zip(ranges_lower) {
438 if R::is_zero(&a[r]) {
439 continue;
440 }
441 let e = a[r].clone();
442 a[r] = R::zero();
443 R::add_scaled_assign(&mut a[(r + 1)..], &pivot_a[(r + 1)..], &R::neg(&e));
444 R::add_scaled_assign(
445 &mut inverse[left..right],
446 &pivot_inverse[left..right],
447 &R::neg(&e),
448 );
449 range.0 = range.0.min(left);
450 range.1 = range.1.max(right);
451 }
452 }
453 for r in (0..n).rev() {
454 let (left, right) = ranges[r];
455 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r);
456 let pivot_inverse = &inverse_lower[0];
457 let (ranges_upper, _) = ranges.split_at_mut(r);
458 for ((a, inverse), range) in a.data[..r].iter_mut().zip(inverse_upper).zip(ranges_upper)
459 {
460 if R::is_zero(&a[r]) {
461 continue;
462 }
463 let e = a[r].clone();
464 a[r] = R::zero();
465 R::add_scaled_assign(
466 &mut inverse[left..right],
467 &pivot_inverse[left..right],
468 &R::neg(&e),
469 );
470 range.0 = range.0.min(left);
471 range.1 = range.1.max(right);
472 }
473 }
474 Some(inverse)
475 }
476
477 pub fn characteristic_polynomial(&mut self) -> Vec<R::T> {
478 let n = self.shape.0;
479 if n == 0 {
480 return vec![R::one()];
481 }
482 assert!(self.data.iter().all(|a| a.len() == n));
483 for j in 0..(n - 1) {
484 if let Some(x) = ((j + 1)..n).find(|&x| !R::is_zero(&self[x][j])) {
485 self.data.swap(j + 1, x);
486 self.data.iter_mut().for_each(|a| a.swap(j + 1, x));
487 let inv = R::inv(&self[j + 1][j]);
488 let mut v = vec![];
489 let src = std::mem::take(&mut self[j + 1]);
490 for a in self.data[(j + 2)..].iter_mut() {
491 let mul = R::mul(&a[j], &inv);
492 R::add_scaled_assign(&mut a[j..], &src[j..], &R::neg(&mul));
493 v.push(mul);
494 }
495 self[j + 1] = src;
496 for a in self.data.iter_mut() {
497 let v = R::dot_product(&a[(j + 2)..], &v);
498 R::add_assign(&mut a[j + 1], &v);
499 }
500 }
501 }
502 // dp[k][j - k] stores [x^k] det(xI - A[..j, ..j]).
503 let mut dp: Vec<Vec<R::T>> = (0..=n).map(|i| Vec::with_capacity(n + 1 - i)).collect();
504 dp[0].push(R::one());
505 for i in 0..n {
506 let mut c = vec![R::zero(); i + 1];
507 c[i] = R::neg(&self[i][i]);
508 let mut mul = R::one();
509 for j in (0..i).rev() {
510 mul = R::mul(&mul, &self[j + 1][j]);
511 c[j] = R::neg(&R::mul(&mul, &self[j][i]));
512 }
513 for k in (0..=i).rev() {
514 let mut value = R::dot_product(&dp[k], &c[k..]);
515 if k > 0 {
516 R::add_assign(&mut value, dp[k - 1].last().unwrap());
517 }
518 dp[k].push(value);
519 }
520 dp[i + 1].push(R::one());
521 }
522 dp.into_iter().map(|mut c| c.pop().unwrap()).collect()
523 }Sourcefn add_assign(x: &mut Self::T, y: &Self::T)
fn add_assign(x: &mut Self::T, y: &Self::T)
Examples found in repository?
crates/competitive/src/algebra/ring.rs (line 48)
45 fn dot_product(x: &[Self::T], y: &[Self::T]) -> Self::T {
46 assert_eq!(x.len(), y.len());
47 x.iter().zip(y).fold(Self::zero(), |mut sum, (x, y)| {
48 Self::add_assign(&mut sum, &Self::mul(x, y));
49 sum
50 })
51 }
52
53 fn add_scaled_assign(x: &mut [Self::T], y: &[Self::T], a: &Self::T) {
54 assert_eq!(x.len(), y.len());
55 for (x, y) in x.iter_mut().zip(y) {
56 Self::add_assign(x, &Self::mul(a, y));
57 }
58 }More examples
crates/competitive/src/math/subset_convolve.rs (line 75)
62 fn transform(t: Self::T, len: usize) -> Self::F {
63 let (mut f, width) = Self::ranked(t, len);
64 let k = width - 1;
65 for bit in 0..k {
66 let half = 1 << bit;
67 for base in (0..len).step_by(half * 2) {
68 for lower in base..base + half {
69 let upper = lower + half;
70 let ranks = lower.count_ones() as usize + 1;
71 let (lower_rows, upper_rows) = f.split_at_mut(upper * width);
72 let lower_row = &lower_rows[lower * width..lower * width + ranks];
73 let upper_row = &mut upper_rows[..ranks];
74 for (upper, lower) in upper_row.iter_mut().zip(lower_row) {
75 R::add_assign(upper, lower);
76 }
77 }
78 }
79 }
80 (f, width)
81 }
82
83 fn inverse_transform((mut f, width): Self::F, len: usize) -> Self::T {
84 let k = width - 1;
85 for bit in 0..k {
86 let half = 1 << bit;
87 for base in (0..len).step_by(half * 2) {
88 for lower in base..base + half {
89 let upper = lower + half;
90 let rank = lower.count_ones() as usize;
91 let (lower_rows, upper_rows) = f.split_at_mut(upper * width);
92 let lower_row = &lower_rows[lower * width + rank..lower * width + width];
93 let upper_row = &mut upper_rows[rank..width];
94 for (upper, lower) in upper_row.iter_mut().zip(lower_row) {
95 R::sub_assign(upper, lower);
96 }
97 }
98 }
99 }
100 Self::diagonal(f, width)
101 }
102
103 fn multiply(f: &mut Self::F, g: &Self::F) {
104 let (f, width) = f;
105 let (g, _) = g;
106 let mut right = vec![R::zero(); *width];
107 let mut output = vec![R::zero(); *width];
108 for (i, f) in f.chunks_exact_mut(*width).enumerate() {
109 let rank = i.count_ones() as usize;
110 let g = &g[i * *width..(i + 1) * *width];
111 let end = Self::multiply_row(f, g, &mut right, &mut output, rank);
112 f[rank..=end].clone_from_slice(&output[rank..=end]);
113 }
114 }
115
116 fn convolve(a: Self::T, b: Self::T) -> Self::T {
117 assert_eq!(a.len(), b.len());
118 let len = a.len();
119 let same = a == b;
120 let (mut x, width) = Self::ranked(a, len);
121 let (mut y, _) = if same {
122 (x.clone(), width)
123 } else {
124 Self::ranked(b, len)
125 };
126 let mut right = vec![R::zero(); width];
127 let mut output = vec![R::zero(); width];
128 for i in 0..len {
129 for bit in (0..(i | len).trailing_zeros() as usize).rev() {
130 let half = width << bit;
131 let start = i * width;
132 let (lower, upper) = x[start..start + half * 2].split_at_mut(half);
133 for (upper, lower) in upper.iter_mut().zip(lower) {
134 R::add_assign(upper, lower);
135 }
136 let (lower, upper) = y[start..start + half * 2].split_at_mut(half);
137 for (upper, lower) in upper.iter_mut().zip(lower) {
138 R::add_assign(upper, lower);
139 }
140 }
141
142 let rank = i.count_ones() as usize;
143 let start = i * width;
144 let x_row = &x[start..start + width];
145 let y_row = &y[start..start + width];
146 output.fill(R::zero());
147 Self::multiply_row(x_row, y_row, &mut right, &mut output, rank);
148 x[start..start + width].clone_from_slice(&output);
149
150 for bit in 0..i.trailing_ones() as usize {
151 let end = (i + 1) * width;
152 let half = width << bit;
153 let (lower, upper) = x[end - half * 2..end].split_at_mut(half);
154 for (upper, lower) in upper.iter_mut().zip(lower) {
155 R::sub_assign(upper, lower);
156 }
157 }
158 }
159 Self::diagonal(x, width)
160 }crates/competitive/src/math/floor_sum.rs (line 268)
213 fn operate(a: &Self::T, b: &Self::T) -> Self::T {
214 let mut a = a.clone();
215 let mut b = b.clone();
216 let mut pow_x = array![R::zero(); X];
217 let mut pow_y = array![R::zero(); Y];
218 pow_x[0] = R::one();
219 pow_y[0] = R::one();
220 for i in 1..X {
221 pow_x[i] = R::mul(&pow_x[i - 1], &a.dx);
222 }
223 for j in 1..Y {
224 pow_y[j] = R::mul(&pow_y[j - 1], &a.dy);
225 }
226 macro_rules! go {
227 ($N:ident) => {
228 let mut comb = array![array![R::zero(); $N]; $N];
229 comb[0][0] = R::one();
230 let mut i = 0;
231 while i + 1 < $N {
232 let mut j = 0;
233 while j <= i {
234 let x = comb[i][j].clone();
235 R::add_assign(&mut comb[i + 1][j], &x);
236 R::add_assign(&mut comb[i + 1][j + 1], &x);
237 j += 1;
238 }
239 i += 1;
240 }
241 for i in 0..X {
242 for j in (0..Y).rev() {
243 for k in j + 1..Y {
244 let mut x = b.dp[i][j].clone();
245 R::mul_assign(&mut x, &comb[k][j]);
246 R::mul_assign(&mut x, &pow_y[k - j]);
247 R::add_assign(&mut b.dp[i][k], &x);
248 }
249 }
250 }
251 for j in 0..Y {
252 for i in (0..X).rev() {
253 for k in i..X {
254 let mut x = b.dp[i][j].clone();
255 R::mul_assign(&mut x, &comb[k][i]);
256 R::mul_assign(&mut x, &pow_x[k - i]);
257 R::add_assign(&mut a.dp[k][j], &x);
258 }
259 }
260 }
261 };
262 }
263 if X <= Y {
264 go!(Y);
265 } else {
266 go!(X);
267 }
268 R::add_assign(&mut a.dx, &b.dx);
269 R::add_assign(&mut a.dy, &b.dy);
270 a
271 }crates/competitive/src/math/matrix.rs (line 235)
158 fn eliminate<const DETERMINANT: bool>(&mut self) -> (usize, R::T) {
159 let (n, m) = self.shape;
160 let mut rank = 0;
161 let mut determinant = R::one();
162 let mut negative = false;
163 for first in (0..m).step_by(64) {
164 if rank == n {
165 break;
166 }
167 let end = (first + 64).min(m);
168 let start = rank;
169 let mut pivots = Vec::new();
170 let mut panel = vec![vec![R::zero(); end - first]; end - first];
171 for col in first..end {
172 if panel[col - first][..col - first]
173 .iter()
174 .any(|x| !R::is_zero(x))
175 {
176 for row in &mut self.data[rank..] {
177 let value =
178 R::dot_product(&row[first..col], &panel[col - first][..col - first]);
179 R::sub_assign(&mut row[col], &value);
180 }
181 }
182 let Some(pivot) = (rank..n).find(|&i| !R::is_zero(&self[i][col])) else {
183 continue;
184 };
185 if pivot != rank {
186 self.data.swap(rank, pivot);
187 negative = !negative;
188 }
189 if DETERMINANT {
190 R::mul_assign(&mut determinant, &self[rank][col]);
191 }
192 let inv = R::inv(&self[rank][col]);
193 let row = &mut self.data[rank];
194 for c in col + 1..end {
195 let value = R::dot_product(&row[first..col], &panel[c - first][..col - first]);
196 R::sub_assign(&mut row[c], &value);
197 panel[c - first][col - first] = row[c].clone();
198 }
199 for row in &mut self.data[rank + 1..] {
200 R::mul_assign(&mut row[col], &inv);
201 }
202 pivots.push(col);
203 rank += 1;
204 if rank == n {
205 break;
206 }
207 }
208 for i in start..if rank - start < 32 { n } else { rank } {
209 let (upper, lower) = self.data.split_at_mut(i);
210 let row = &mut lower[0];
211 for (j, &col) in pivots[..(i - start).min(pivots.len())].iter().enumerate() {
212 if R::is_zero(&row[col]) {
213 continue;
214 }
215 let factor = R::neg(&row[col]);
216 R::add_scaled_assign(&mut row[end..], &upper[start + j][end..], &factor);
217 }
218 }
219 if rank < n
220 && end < m
221 && rank - start >= 32
222 && self.data[rank..]
223 .iter()
224 .any(|row| pivots.iter().any(|&col| !R::is_zero(&row[col])))
225 {
226 let lower = Self::new_with((n - rank, rank - start), |i, j| {
227 R::neg(&self[rank + i][pivots[j]])
228 });
229 let upper = Self::new_with((rank - start, m - end), |i, j| {
230 self[start + i][end + j].clone()
231 });
232 let update = &lower * &upper;
233 for (row, update) in self.data[rank..].iter_mut().zip(update.data) {
234 for (x, y) in row[end..].iter_mut().zip(update) {
235 R::add_assign(x, &y);
236 }
237 }
238 }
239 for (i, &col) in pivots.iter().enumerate() {
240 for row in &mut self.data[start + i + 1..] {
241 row[col] = R::zero();
242 }
243 }
244 if DETERMINANT && rank < end {
245 return (rank, R::zero());
246 }
247 }
248 if DETERMINANT && negative {
249 determinant = R::neg(&determinant);
250 }
251 (rank, determinant)
252 }
253
254 /// f: (row, pivot_row, col)
255 pub fn row_reduction_with<F>(&mut self, normalize: bool, mut f: F)
256 where
257 F: FnMut(usize, usize, usize),
258 {
259 let (n, m) = self.shape;
260 let mut c = 0;
261 for r in 0..n {
262 loop {
263 if c >= m {
264 return;
265 }
266 if let Some(pivot) = (r..n).find(|&p| !R::is_zero(&self[p][c])) {
267 f(r, pivot, c);
268 self.data.swap(r, pivot);
269 break;
270 };
271 c += 1;
272 }
273 let d = R::inv(&self[r][c]);
274 if normalize {
275 for value in &mut self[r][c..m] {
276 R::mul_assign(value, &d);
277 }
278 }
279 for i in (0..n).filter(|&i| i != r) {
280 let mut e = self[i][c].clone();
281 if !normalize {
282 R::mul_assign(&mut e, &d);
283 }
284 for j in c..m {
285 let e = R::mul(&e, &self[r][j]);
286 R::sub_assign(&mut self[i][j], &e);
287 }
288 }
289 c += 1;
290 }
291 }
292
293 pub fn row_reduction(&mut self, normalize: bool) {
294 self.row_reduction_with(normalize, |_, _, _| {});
295 }
296
297 pub fn rank(&mut self) -> usize {
298 self.eliminate::<false>().0
299 }
300
301 pub fn determinant(&mut self) -> R::T {
302 assert_eq!(self.shape.0, self.shape.1);
303 self.eliminate::<true>().1
304 }
305
306 pub fn solve_system_of_linear_equations(
307 &self,
308 b: &[R::T],
309 ) -> Option<SystemOfLinearEquationsSolution<R>> {
310 assert_eq!(self.shape.0, b.len());
311 let m = self.shape.1;
312 let mut a = Self::new_with((self.shape.0, m + 1), |i, j| {
313 if j == m {
314 b[i].clone()
315 } else {
316 self[i][j].clone()
317 }
318 });
319 let rank = a.eliminate::<false>().0;
320 let mut pivots = Vec::with_capacity(rank);
321 let mut b = Vec::with_capacity(rank);
322 for row in &a.data[..rank] {
323 let c = row.iter().position(|x| !R::is_zero(x)).unwrap();
324 if c == m {
325 return None;
326 }
327 pivots.push(c);
328 b.push(row[m].clone());
329 }
330
331 let mut free = Vec::with_capacity(m - rank);
332 let mut pivot = 0;
333 for c in 0..m {
334 if pivot < rank && pivots[pivot] == c {
335 pivot += 1;
336 } else {
337 free.push(c);
338 }
339 }
340 let mut coefficients: Vec<Vec<_>> = (0..rank)
341 .map(|i| free.iter().map(|&c| a[i][c].clone()).collect())
342 .collect();
343 for k in (0..rank).rev() {
344 let c = pivots[k];
345 let inv = R::inv(&a[k][c]);
346 R::mul_assign(&mut b[k], &inv);
347 let pivot_b = b[k].clone();
348 let (upper, lower) = coefficients.split_at_mut(k);
349 let pivot_coefficients = &mut lower[0];
350 for x in pivot_coefficients.iter_mut() {
351 R::mul_assign(x, &inv);
352 }
353 for ((row, value), coefficients) in a.data[..k].iter_mut().zip(&mut b[..k]).zip(upper) {
354 if R::is_zero(&row[c]) {
355 continue;
356 }
357 let factor = row[c].clone();
358 row[c] = R::zero();
359 R::sub_assign(value, &R::mul(&factor, &pivot_b));
360 R::add_scaled_assign(coefficients, pivot_coefficients, &R::neg(&factor));
361 }
362 }
363
364 let mut particular = vec![R::zero(); m];
365 for i in 0..rank {
366 particular[pivots[i]] = b[i].clone();
367 }
368 let mut basis = Vec::with_capacity(free.len());
369 for (j, &c) in free.iter().enumerate() {
370 let mut vector = vec![R::zero(); m];
371 vector[c] = R::one();
372 for i in 0..rank {
373 vector[pivots[i]] = R::neg(&coefficients[i][j]);
374 }
375 basis.push(vector);
376 }
377 Some(SystemOfLinearEquationsSolution { particular, basis })
378 }
379
380 pub fn inverse(&self) -> Option<Matrix<R>> {
381 assert_eq!(self.shape.0, self.shape.1);
382 let n = self.shape.0;
383 if n >= 64 {
384 let m = n / 2;
385 let a = Self::new_with((m, m), |i, j| self[i][j].clone());
386 if let Some(mut ai) = a.inverse() {
387 let b = Self::new_with((m, n - m), |i, j| self[i][j + m].clone());
388 let c = Self::new_with((n - m, m), |i, j| self[i + m][j].clone());
389 let mut d = Self::new_with((n - m, n - m), |i, j| self[i + m][j + m].clone());
390 let u = &ai * &b;
391 let v = &c * &ai;
392 d -= &v * &b;
393 let di = d.inverse()?;
394 let r = &u * &di;
395 let t = &di * &v;
396 ai += &r * &v;
397 let mut inverse = Self::zeros((n, n));
398 for i in 0..m {
399 inverse[i][..m].clone_from_slice(&ai[i]);
400 for (x, y) in inverse[i][m..].iter_mut().zip(&r[i]) {
401 *x = R::neg(y);
402 }
403 }
404 for i in m..n {
405 for (x, y) in inverse[i][..m].iter_mut().zip(&t[i - m]) {
406 *x = R::neg(y);
407 }
408 inverse[i][m..].clone_from_slice(&di[i - m]);
409 }
410 return Some(inverse);
411 }
412 }
413 let mut a = self.clone();
414 let mut inverse = Self::eye((n, n));
415 let mut ranges: Vec<_> = (0..n).map(|i| (i, i + 1)).collect();
416 for r in 0..n {
417 let pivot = (r..n).find(|&i| !R::is_zero(&a[i][r]))?;
418 a.data.swap(r, pivot);
419 inverse.data.swap(r, pivot);
420 ranges.swap(r, pivot);
421
422 let d = R::inv(&a[r][r]);
423 for x in &mut a[r][r..] {
424 R::mul_assign(x, &d);
425 }
426 let (left, right) = ranges[r];
427 for x in &mut inverse[r][left..right] {
428 R::mul_assign(x, &d);
429 }
430
431 let (a_upper, a_lower) = a.data.split_at_mut(r + 1);
432 let pivot_a = &a_upper[r];
433 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r + 1);
434 let pivot_inverse = &inverse_upper[r];
435 let (ranges_upper, ranges_lower) = ranges.split_at_mut(r + 1);
436 let (left, right) = ranges_upper[r];
437 for ((a, inverse), range) in a_lower.iter_mut().zip(inverse_lower).zip(ranges_lower) {
438 if R::is_zero(&a[r]) {
439 continue;
440 }
441 let e = a[r].clone();
442 a[r] = R::zero();
443 R::add_scaled_assign(&mut a[(r + 1)..], &pivot_a[(r + 1)..], &R::neg(&e));
444 R::add_scaled_assign(
445 &mut inverse[left..right],
446 &pivot_inverse[left..right],
447 &R::neg(&e),
448 );
449 range.0 = range.0.min(left);
450 range.1 = range.1.max(right);
451 }
452 }
453 for r in (0..n).rev() {
454 let (left, right) = ranges[r];
455 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r);
456 let pivot_inverse = &inverse_lower[0];
457 let (ranges_upper, _) = ranges.split_at_mut(r);
458 for ((a, inverse), range) in a.data[..r].iter_mut().zip(inverse_upper).zip(ranges_upper)
459 {
460 if R::is_zero(&a[r]) {
461 continue;
462 }
463 let e = a[r].clone();
464 a[r] = R::zero();
465 R::add_scaled_assign(
466 &mut inverse[left..right],
467 &pivot_inverse[left..right],
468 &R::neg(&e),
469 );
470 range.0 = range.0.min(left);
471 range.1 = range.1.max(right);
472 }
473 }
474 Some(inverse)
475 }
476
477 pub fn characteristic_polynomial(&mut self) -> Vec<R::T> {
478 let n = self.shape.0;
479 if n == 0 {
480 return vec![R::one()];
481 }
482 assert!(self.data.iter().all(|a| a.len() == n));
483 for j in 0..(n - 1) {
484 if let Some(x) = ((j + 1)..n).find(|&x| !R::is_zero(&self[x][j])) {
485 self.data.swap(j + 1, x);
486 self.data.iter_mut().for_each(|a| a.swap(j + 1, x));
487 let inv = R::inv(&self[j + 1][j]);
488 let mut v = vec![];
489 let src = std::mem::take(&mut self[j + 1]);
490 for a in self.data[(j + 2)..].iter_mut() {
491 let mul = R::mul(&a[j], &inv);
492 R::add_scaled_assign(&mut a[j..], &src[j..], &R::neg(&mul));
493 v.push(mul);
494 }
495 self[j + 1] = src;
496 for a in self.data.iter_mut() {
497 let v = R::dot_product(&a[(j + 2)..], &v);
498 R::add_assign(&mut a[j + 1], &v);
499 }
500 }
501 }
502 // dp[k][j - k] stores [x^k] det(xI - A[..j, ..j]).
503 let mut dp: Vec<Vec<R::T>> = (0..=n).map(|i| Vec::with_capacity(n + 1 - i)).collect();
504 dp[0].push(R::one());
505 for i in 0..n {
506 let mut c = vec![R::zero(); i + 1];
507 c[i] = R::neg(&self[i][i]);
508 let mut mul = R::one();
509 for j in (0..i).rev() {
510 mul = R::mul(&mul, &self[j + 1][j]);
511 c[j] = R::neg(&R::mul(&mul, &self[j][i]));
512 }
513 for k in (0..=i).rev() {
514 let mut value = R::dot_product(&dp[k], &c[k..]);
515 if k > 0 {
516 R::add_assign(&mut value, dp[k - 1].last().unwrap());
517 }
518 dp[k].push(value);
519 }
520 dp[i + 1].push(R::one());
521 }
522 dp.into_iter().map(|mut c| c.pop().unwrap()).collect()
523 }
524}
525
526impl<R> Index<usize> for Matrix<R>
527where
528 R: SemiRing,
529{
530 type Output = Vec<R::T>;
531 fn index(&self, index: usize) -> &Self::Output {
532 &self.data[index]
533 }
534}
535
536impl<R> IndexMut<usize> for Matrix<R>
537where
538 R: SemiRing,
539{
540 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
541 &mut self.data[index]
542 }
543}
544
545impl<R> Index<(usize, usize)> for Matrix<R>
546where
547 R: SemiRing,
548{
549 type Output = R::T;
550 fn index(&self, index: (usize, usize)) -> &Self::Output {
551 &self.data[index.0][index.1]
552 }
553}
554
555impl<R> IndexMut<(usize, usize)> for Matrix<R>
556where
557 R: SemiRing,
558{
559 fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
560 &mut self.data[index.0][index.1]
561 }
562}
563
564macro_rules! impl_matrix_pairwise_binop {
565 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident $(where [$($clauses:tt)*])?) => {
566 impl<R> $imp_assign for Matrix<R>
567 where
568 R: SemiRing,
569 $($($clauses)*)?
570 {
571 fn $method_assign(&mut self, rhs: Self) {
572 self.pairwise_assign(&rhs, |a, b| R::$method_assign(a, b));
573 }
574 }
575 impl<R> $imp_assign<&Matrix<R>> for Matrix<R>
576 where
577 R: SemiRing,
578 $($($clauses)*)?
579 {
580 fn $method_assign(&mut self, rhs: &Self) {
581 self.pairwise_assign(rhs, |a, b| R::$method_assign(a, b));
582 }
583 }
584 impl<R> $imp for Matrix<R>
585 where
586 R: SemiRing,
587 $($($clauses)*)?
588 {
589 type Output = Matrix<R>;
590 fn $method(mut self, rhs: Self) -> Self::Output {
591 self.$method_assign(rhs);
592 self
593 }
594 }
595 impl<R> $imp<&Matrix<R>> for Matrix<R>
596 where
597 R: SemiRing,
598 $($($clauses)*)?
599 {
600 type Output = Matrix<R>;
601 fn $method(mut self, rhs: &Self) -> Self::Output {
602 self.$method_assign(rhs);
603 self
604 }
605 }
606 impl<R> $imp<Matrix<R>> for &Matrix<R>
607 where
608 R: SemiRing,
609 $($($clauses)*)?
610 {
611 type Output = Matrix<R>;
612 fn $method(self, mut rhs: Matrix<R>) -> Self::Output {
613 rhs.pairwise_assign(self, |a, b| *a = R::$method(b, a));
614 rhs
615 }
616 }
617 impl<R> $imp<&Matrix<R>> for &Matrix<R>
618 where
619 R: SemiRing,
620 $($($clauses)*)?
621 {
622 type Output = Matrix<R>;
623 fn $method(self, rhs: &Matrix<R>) -> Self::Output {
624 let mut this = self.clone();
625 this.$method_assign(rhs);
626 this
627 }
628 }
629 };
630}
631
632impl_matrix_pairwise_binop!(Add, add, AddAssign, add_assign);
633impl_matrix_pairwise_binop!(Sub, sub, SubAssign, sub_assign where [R: SemiRing<Additive: Invertible>]);
634
635impl<R> Mul for Matrix<R>
636where
637 R: SemiRing,
638{
639 type Output = Matrix<R>;
640 fn mul(self, rhs: Self) -> Self::Output {
641 (&self).mul(&rhs)
642 }
643}
644impl<R> Mul<&Matrix<R>> for Matrix<R>
645where
646 R: SemiRing,
647{
648 type Output = Matrix<R>;
649 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
650 (&self).mul(rhs)
651 }
652}
653impl<R> Mul<Matrix<R>> for &Matrix<R>
654where
655 R: SemiRing,
656{
657 type Output = Matrix<R>;
658 fn mul(self, rhs: Matrix<R>) -> Self::Output {
659 self.mul(&rhs)
660 }
661}
662impl<R> Mul<&Matrix<R>> for &Matrix<R>
663where
664 R: SemiRing,
665{
666 type Output = Matrix<R>;
667 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
668 assert_eq!(self.shape.1, rhs.shape.0);
669 if let Some(data) = R::try_matrix_product(&self.data, &rhs.data) {
670 return Matrix::from_vec(data);
671 }
672 let rhs = rhs.transpose();
673 Matrix::new_with((self.shape.0, rhs.shape.0), |i, j| {
674 R::dot_product(&self[i], &rhs[j])
675 })
676 }
677}
678
679fn strassen_rec<R: Ring>(
680 a: &[R::T],
681 b: &[R::T],
682 c: &mut [R::T],
683 shape: (usize, usize, usize),
684 stride_a: usize,
685 stride_b: usize,
686) {
687 let (n, m, p) = shape;
688 fn add_block<R: Ring>(
689 a: &[R::T],
690 b: &[R::T],
691 out: &mut [R::T],
692 n: usize,
693 stride_a: usize,
694 stride_b: usize,
695 ) {
696 for ((a, b), c) in a
697 .chunks(stride_a)
698 .zip(b.chunks(stride_b))
699 .zip(out.chunks_exact_mut(n))
700 {
701 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
702 *c = R::add(a, b);
703 }
704 }
705 }
706
707 fn sub_block<R: Ring>(
708 a: &[R::T],
709 b: &[R::T],
710 out: &mut [R::T],
711 n: usize,
712 stride_a: usize,
713 stride_b: usize,
714 ) {
715 for ((a, b), c) in a
716 .chunks(stride_a)
717 .zip(b.chunks(stride_b))
718 .zip(out.chunks_exact_mut(n))
719 {
720 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
721 *c = R::sub(a, b);
722 }
723 }
724 }
725
726 if n.min(m).min(p) <= 128 {
727 let transposed: Vec<_> = (0..p)
728 .flat_map(|j| (0..m).map(move |i| b[i * stride_b + j].clone()))
729 .collect();
730 for (a, c) in a.chunks(stride_a).zip(c.chunks_exact_mut(p)) {
731 for (b, c) in transposed.chunks_exact(m).zip(c) {
732 *c = R::dot_product(&a[..m], b);
733 }
734 }
735 return;
736 }
737 let (h, k, w) = (n / 2, m / 2, p / 2);
738 let a11 = 0;
739 let a12 = k;
740 let a21 = h * stride_a;
741 let a22 = a21 + k;
742 let b11 = 0;
743 let b12 = w;
744 let b21 = k * stride_b;
745 let b22 = b21 + w;
746
747 let block = h * w;
748 let mut buf = vec![R::zero(); h * k + k * w + block * 7];
749 let (s1, rest) = buf.split_at_mut(h * k);
750 let (s2, m_buf) = rest.split_at_mut(k * w);
751 let (m1, rest) = m_buf.split_at_mut(block);
752 let (m2, rest) = rest.split_at_mut(block);
753 let (m3, rest) = rest.split_at_mut(block);
754 let (m4, rest) = rest.split_at_mut(block);
755 let (m5, rest) = rest.split_at_mut(block);
756 let (m6, m7) = rest.split_at_mut(block);
757
758 // (A11 + A22)(B11 + B22)
759 add_block::<R>(&a[a11..], &a[a22..], s1, k, stride_a, stride_a);
760 add_block::<R>(&b[b11..], &b[b22..], s2, w, stride_b, stride_b);
761 strassen_rec::<R>(s1, s2, m1, (h, k, w), k, w);
762
763 // (A21 + A22) B11
764 add_block::<R>(&a[a21..], &a[a22..], s1, k, stride_a, stride_a);
765 strassen_rec::<R>(s1, &b[b11..], m2, (h, k, w), k, stride_b);
766
767 // A11 (B12 - B22)
768 sub_block::<R>(&b[b12..], &b[b22..], s2, w, stride_b, stride_b);
769 strassen_rec::<R>(&a[a11..], s2, m3, (h, k, w), stride_a, w);
770
771 // A22 (B21 - B11)
772 sub_block::<R>(&b[b21..], &b[b11..], s2, w, stride_b, stride_b);
773 strassen_rec::<R>(&a[a22..], s2, m4, (h, k, w), stride_a, w);
774
775 // (A11 + A12) B22
776 add_block::<R>(&a[a11..], &a[a12..], s1, k, stride_a, stride_a);
777 strassen_rec::<R>(s1, &b[b22..], m5, (h, k, w), k, stride_b);
778
779 // (A21 - A11)(B11 + B12)
780 sub_block::<R>(&a[a21..], &a[a11..], s1, k, stride_a, stride_a);
781 add_block::<R>(&b[b11..], &b[b12..], s2, w, stride_b, stride_b);
782 strassen_rec::<R>(s1, s2, m6, (h, k, w), k, w);
783
784 // (A12 - A22)(B21 + B22)
785 sub_block::<R>(&a[a12..], &a[a22..], s1, k, stride_a, stride_a);
786 add_block::<R>(&b[b21..], &b[b22..], s2, w, stride_b, stride_b);
787 strassen_rec::<R>(s1, s2, m7, (h, k, w), k, w);
788
789 let c11 = 0;
790 let c12 = w;
791 let c21 = h * p;
792 let c22 = c21 + w;
793 for ((((m1, m4), m5), m7), c) in m1
794 .iter()
795 .zip(m4.iter())
796 .zip(m5.iter())
797 .zip(m7.iter())
798 .zip(c[c11..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
799 {
800 *c = R::add(m1, m4);
801 R::sub_assign(c, m5);
802 R::add_assign(c, m7);
803 }
804 for ((m3, m5), c) in m3
805 .iter()
806 .zip(m5.iter())
807 .zip(c[c12..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
808 {
809 *c = R::add(m3, m5);
810 }
811 for ((m2, m4), c) in m2
812 .iter()
813 .zip(m4.iter())
814 .zip(c[c21..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
815 {
816 *c = R::add(m2, m4);
817 }
818 for ((((m1, m2), m3), m6), c) in m1
819 .iter()
820 .zip(m2.iter())
821 .zip(m3.iter())
822 .zip(m6.iter())
823 .zip(c[c22..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
824 {
825 *c = R::sub(m1, m2);
826 R::add_assign(c, m3);
827 R::add_assign(c, m6);
828 }
829}crates/competitive/src/math/black_box_matrix.rs (line 110)
97 pub fn determinant(&self) -> R::T {
98 assert_eq!(self.shape.0, self.shape.1);
99 let n = self.shape.0;
100 let mut columns = vec![Vec::<(usize, R::T)>::new(); n];
101 for &(i, j, ref value) in &self.nonzero {
102 columns[j].push((i, value.clone()));
103 }
104 let mut degrees = vec![0; n];
105 for column in &mut columns {
106 column.sort_unstable_by_key(|&(i, _)| i);
107 let mut merged: Vec<(usize, R::T)> = Vec::with_capacity(column.len());
108 for (i, value) in column.drain(..) {
109 if let Some((_, x)) = merged.last_mut().filter(|(last, _)| *last == i) {
110 R::add_assign(x, &value);
111 } else {
112 merged.push((i, value));
113 }
114 }
115 merged.retain(|(i, value)| {
116 if R::is_zero(value) {
117 false
118 } else {
119 degrees[*i] += 1;
120 true
121 }
122 });
123 *column = merged;
124 }
125 let mut order: Vec<_> = (0..n).collect();
126 order.sort_unstable_by_key(|&j| columns[j].len());
127 let mut lower: Vec<Vec<(usize, R::T)>> = Vec::with_capacity(n);
128 let mut pivots: Vec<Option<usize>> = vec![None; n];
129 let mut x = vec![R::zero(); n];
130 let mut seen = vec![0; n];
131 let mut stack = Vec::new();
132 let mut support = Vec::new();
133 let mut determinant = R::one();
134 for (k, &j) in order.iter().enumerate() {
135 support.clear();
136 for &(i, _) in &columns[j] {
137 if seen[i] == k + 1 {
138 continue;
139 }
140 seen[i] = k + 1;
141 x[i] = R::zero();
142 stack.push((i, 0));
143 while let Some((i, next)) = stack.last_mut() {
144 if let Some(pivot) = pivots[*i]
145 && *next < lower[pivot].len()
146 {
147 let row = lower[pivot][*next].0;
148 *next += 1;
149 if seen[row] != k + 1 {
150 seen[row] = k + 1;
151 x[row] = R::zero();
152 stack.push((row, 0));
153 }
154 continue;
155 }
156 support.push(*i);
157 stack.pop();
158 }
159 }
160 for &(i, ref value) in &columns[j] {
161 x[i] = value.clone();
162 }
163 let mut pivot = None;
164 for &i in support.iter().rev() {
165 if let Some(p) = pivots[i] {
166 let factor = R::neg(&x[i]);
167 for &(row, ref value) in &lower[p] {
168 R::add_assign(&mut x[row], &R::mul(&factor, value));
169 }
170 } else if !R::is_zero(&x[i]) && pivot.is_none_or(|p| degrees[i] < degrees[p]) {
171 pivot = Some(i);
172 }
173 }
174 let Some(pivot) = pivot else { return R::zero() };
175 R::mul_assign(&mut determinant, &x[pivot]);
176 let inv = R::inv(&x[pivot]);
177 pivots[pivot] = Some(k);
178 lower.push(
179 support
180 .iter()
181 .filter(|&&i| pivots[i].is_none() && !R::is_zero(&x[i]))
182 .map(|&i| (i, R::mul(&x[i], &inv)))
183 .collect(),
184 );
185 }
186 for mut permutation in [order, pivots.into_iter().map(Option::unwrap).collect()] {
187 for i in 0..n {
188 while permutation[i] != i {
189 let j = permutation[i];
190 permutation.swap(i, j);
191 determinant = R::neg(&determinant);
192 }
193 }
194 }
195 determinant
196 }
197}
198
199impl<R> From<Matrix<R>> for SparseMatrix<R>
200where
201 R: SemiRing<T: PartialEq>,
202{
203 fn from(mat: Matrix<R>) -> Self {
204 let mut nonzero = vec![];
205 for i in 0..mat.shape.0 {
206 for j in 0..mat.shape.1 {
207 let v = mat[(i, j)].clone();
208 if !R::is_zero(&v) {
209 nonzero.push((i, j, v));
210 }
211 }
212 }
213 Self {
214 shape: mat.shape,
215 nonzero,
216 }
217 }
218}
219
220impl<R> From<SparseMatrix<R>> for Matrix<R>
221where
222 R: SemiRing,
223{
224 fn from(smat: SparseMatrix<R>) -> Self {
225 let mut mat = Matrix::zeros(smat.shape);
226 for &(i, j, ref v) in &smat.nonzero {
227 R::add_assign(&mut mat[(i, j)], v);
228 }
229 mat
230 }
231}
232
233impl<R> BlackBoxMatrix<R> for SparseMatrix<R>
234where
235 R: SemiRing,
236{
237 fn apply(&self, v: &[R::T]) -> Vec<R::T> {
238 assert_eq!(self.shape.1, v.len());
239 let mut res = vec![R::zero(); self.shape.0];
240 for &(i, j, ref val) in &self.nonzero {
241 R::add_assign(&mut res[i], &R::mul(val, &v[j]));
242 }
243 res
244 }crates/competitive/src/algorithm/automata_learning.rs (line 558)
523 pub fn train_sample(&mut self, sample: &[usize]) -> bool {
524 let Some((prefix, suffix)) = self.split_sample(sample) else {
525 return false;
526 };
527 self.prefixes.push(prefix);
528 self.suffixes.push(suffix);
529 let n = self.inv_h.shape.0;
530 let prefix = &self.prefixes[n];
531 let suffix = &self.suffixes[n];
532 let u = Matrix::<F>::new_with((n, 1), |i, _| {
533 self.automaton.behavior(
534 self.prefixes[i]
535 .iter()
536 .cloned()
537 .chain(suffix.iter().cloned()),
538 )
539 });
540 let v = Matrix::<F>::new_with((1, n), |_, j| {
541 self.automaton.behavior(
542 prefix
543 .iter()
544 .cloned()
545 .chain(self.suffixes[j].iter().cloned()),
546 )
547 });
548 let w = Matrix::<F>::new_with((1, 1), |_, _| {
549 self.automaton
550 .behavior(prefix.iter().cloned().chain(suffix.iter().cloned()))
551 });
552 let t = &self.inv_h * &u;
553 let s = &v * &self.inv_h;
554 let d = F::inv(&(&w - &(&v * &t))[0][0]);
555 let dh = &t * &s;
556 for i in 0..n {
557 for j in 0..n {
558 F::add_assign(&mut self.inv_h[i][j], &F::mul(&dh[i][j], &d));
559 }
560 }
561 self.inv_h
562 .add_col_with(|i, _| F::neg(&F::mul(&t[i][0], &d)));
563 self.inv_h.add_row_with(|_, j| {
564 if j != n {
565 F::neg(&F::mul(&s[0][j], &d))
566 } else {
567 d.clone()
568 }
569 });
570
571 for (x, transition) in self.wfa.transitions.iter_mut().enumerate() {
572 let b = &(&self.nh[x] * &t) * &s;
573 for i in 0..n {
574 for j in 0..n {
575 F::add_assign(&mut transition[i][j], &F::mul(&b[i][j], &d));
576 }
577 }
578 }
579 for (x, nh) in self.nh.iter_mut().enumerate() {
580 nh.add_col_with(|i, j| {
581 self.automaton.behavior(
582 self.prefixes[i]
583 .iter()
584 .cloned()
585 .chain([x])
586 .chain(self.suffixes[j].iter().cloned()),
587 )
588 });
589 nh.add_row_with(|i, j| {
590 self.automaton.behavior(
591 self.prefixes[i]
592 .iter()
593 .cloned()
594 .chain([x])
595 .chain(self.suffixes[j].iter().cloned()),
596 )
597 });
598 }
599 self.wfa
600 .initial_weights
601 .add_col_with(|_, _| if n == 0 { F::one() } else { F::zero() });
602 self.wfa
603 .final_weights
604 .add_row_with(|_, _| self.automaton.behavior(prefix.iter().cloned()));
605 for (x, transition) in self.wfa.transitions.iter_mut().enumerate() {
606 transition.add_col_with(|_, _| F::zero());
607 transition.add_row_with(|_, _| F::zero());
608 for i in 0..=n {
609 for j in 0..=n {
610 if i == n || j == n {
611 for k in 0..=n {
612 if i != n && j != n && k != n {
613 continue;
614 }
615 F::add_assign(
616 &mut transition[i][k],
617 &F::mul(&self.nh[x][i][j], &self.inv_h[j][k]),
618 );
619 }
620 } else {
621 let k = n;
622 F::add_assign(
623 &mut transition[i][k],
624 &F::mul(&self.nh[x][i][j], &self.inv_h[j][k]),
625 );
626 }
627 }
628 }
629 }
630 true
631 }Sourcefn mul_assign(x: &mut Self::T, y: &Self::T)
fn mul_assign(x: &mut Self::T, y: &Self::T)
Examples found in repository?
crates/competitive/src/math/matrix.rs (line 190)
158 fn eliminate<const DETERMINANT: bool>(&mut self) -> (usize, R::T) {
159 let (n, m) = self.shape;
160 let mut rank = 0;
161 let mut determinant = R::one();
162 let mut negative = false;
163 for first in (0..m).step_by(64) {
164 if rank == n {
165 break;
166 }
167 let end = (first + 64).min(m);
168 let start = rank;
169 let mut pivots = Vec::new();
170 let mut panel = vec![vec![R::zero(); end - first]; end - first];
171 for col in first..end {
172 if panel[col - first][..col - first]
173 .iter()
174 .any(|x| !R::is_zero(x))
175 {
176 for row in &mut self.data[rank..] {
177 let value =
178 R::dot_product(&row[first..col], &panel[col - first][..col - first]);
179 R::sub_assign(&mut row[col], &value);
180 }
181 }
182 let Some(pivot) = (rank..n).find(|&i| !R::is_zero(&self[i][col])) else {
183 continue;
184 };
185 if pivot != rank {
186 self.data.swap(rank, pivot);
187 negative = !negative;
188 }
189 if DETERMINANT {
190 R::mul_assign(&mut determinant, &self[rank][col]);
191 }
192 let inv = R::inv(&self[rank][col]);
193 let row = &mut self.data[rank];
194 for c in col + 1..end {
195 let value = R::dot_product(&row[first..col], &panel[c - first][..col - first]);
196 R::sub_assign(&mut row[c], &value);
197 panel[c - first][col - first] = row[c].clone();
198 }
199 for row in &mut self.data[rank + 1..] {
200 R::mul_assign(&mut row[col], &inv);
201 }
202 pivots.push(col);
203 rank += 1;
204 if rank == n {
205 break;
206 }
207 }
208 for i in start..if rank - start < 32 { n } else { rank } {
209 let (upper, lower) = self.data.split_at_mut(i);
210 let row = &mut lower[0];
211 for (j, &col) in pivots[..(i - start).min(pivots.len())].iter().enumerate() {
212 if R::is_zero(&row[col]) {
213 continue;
214 }
215 let factor = R::neg(&row[col]);
216 R::add_scaled_assign(&mut row[end..], &upper[start + j][end..], &factor);
217 }
218 }
219 if rank < n
220 && end < m
221 && rank - start >= 32
222 && self.data[rank..]
223 .iter()
224 .any(|row| pivots.iter().any(|&col| !R::is_zero(&row[col])))
225 {
226 let lower = Self::new_with((n - rank, rank - start), |i, j| {
227 R::neg(&self[rank + i][pivots[j]])
228 });
229 let upper = Self::new_with((rank - start, m - end), |i, j| {
230 self[start + i][end + j].clone()
231 });
232 let update = &lower * &upper;
233 for (row, update) in self.data[rank..].iter_mut().zip(update.data) {
234 for (x, y) in row[end..].iter_mut().zip(update) {
235 R::add_assign(x, &y);
236 }
237 }
238 }
239 for (i, &col) in pivots.iter().enumerate() {
240 for row in &mut self.data[start + i + 1..] {
241 row[col] = R::zero();
242 }
243 }
244 if DETERMINANT && rank < end {
245 return (rank, R::zero());
246 }
247 }
248 if DETERMINANT && negative {
249 determinant = R::neg(&determinant);
250 }
251 (rank, determinant)
252 }
253
254 /// f: (row, pivot_row, col)
255 pub fn row_reduction_with<F>(&mut self, normalize: bool, mut f: F)
256 where
257 F: FnMut(usize, usize, usize),
258 {
259 let (n, m) = self.shape;
260 let mut c = 0;
261 for r in 0..n {
262 loop {
263 if c >= m {
264 return;
265 }
266 if let Some(pivot) = (r..n).find(|&p| !R::is_zero(&self[p][c])) {
267 f(r, pivot, c);
268 self.data.swap(r, pivot);
269 break;
270 };
271 c += 1;
272 }
273 let d = R::inv(&self[r][c]);
274 if normalize {
275 for value in &mut self[r][c..m] {
276 R::mul_assign(value, &d);
277 }
278 }
279 for i in (0..n).filter(|&i| i != r) {
280 let mut e = self[i][c].clone();
281 if !normalize {
282 R::mul_assign(&mut e, &d);
283 }
284 for j in c..m {
285 let e = R::mul(&e, &self[r][j]);
286 R::sub_assign(&mut self[i][j], &e);
287 }
288 }
289 c += 1;
290 }
291 }
292
293 pub fn row_reduction(&mut self, normalize: bool) {
294 self.row_reduction_with(normalize, |_, _, _| {});
295 }
296
297 pub fn rank(&mut self) -> usize {
298 self.eliminate::<false>().0
299 }
300
301 pub fn determinant(&mut self) -> R::T {
302 assert_eq!(self.shape.0, self.shape.1);
303 self.eliminate::<true>().1
304 }
305
306 pub fn solve_system_of_linear_equations(
307 &self,
308 b: &[R::T],
309 ) -> Option<SystemOfLinearEquationsSolution<R>> {
310 assert_eq!(self.shape.0, b.len());
311 let m = self.shape.1;
312 let mut a = Self::new_with((self.shape.0, m + 1), |i, j| {
313 if j == m {
314 b[i].clone()
315 } else {
316 self[i][j].clone()
317 }
318 });
319 let rank = a.eliminate::<false>().0;
320 let mut pivots = Vec::with_capacity(rank);
321 let mut b = Vec::with_capacity(rank);
322 for row in &a.data[..rank] {
323 let c = row.iter().position(|x| !R::is_zero(x)).unwrap();
324 if c == m {
325 return None;
326 }
327 pivots.push(c);
328 b.push(row[m].clone());
329 }
330
331 let mut free = Vec::with_capacity(m - rank);
332 let mut pivot = 0;
333 for c in 0..m {
334 if pivot < rank && pivots[pivot] == c {
335 pivot += 1;
336 } else {
337 free.push(c);
338 }
339 }
340 let mut coefficients: Vec<Vec<_>> = (0..rank)
341 .map(|i| free.iter().map(|&c| a[i][c].clone()).collect())
342 .collect();
343 for k in (0..rank).rev() {
344 let c = pivots[k];
345 let inv = R::inv(&a[k][c]);
346 R::mul_assign(&mut b[k], &inv);
347 let pivot_b = b[k].clone();
348 let (upper, lower) = coefficients.split_at_mut(k);
349 let pivot_coefficients = &mut lower[0];
350 for x in pivot_coefficients.iter_mut() {
351 R::mul_assign(x, &inv);
352 }
353 for ((row, value), coefficients) in a.data[..k].iter_mut().zip(&mut b[..k]).zip(upper) {
354 if R::is_zero(&row[c]) {
355 continue;
356 }
357 let factor = row[c].clone();
358 row[c] = R::zero();
359 R::sub_assign(value, &R::mul(&factor, &pivot_b));
360 R::add_scaled_assign(coefficients, pivot_coefficients, &R::neg(&factor));
361 }
362 }
363
364 let mut particular = vec![R::zero(); m];
365 for i in 0..rank {
366 particular[pivots[i]] = b[i].clone();
367 }
368 let mut basis = Vec::with_capacity(free.len());
369 for (j, &c) in free.iter().enumerate() {
370 let mut vector = vec![R::zero(); m];
371 vector[c] = R::one();
372 for i in 0..rank {
373 vector[pivots[i]] = R::neg(&coefficients[i][j]);
374 }
375 basis.push(vector);
376 }
377 Some(SystemOfLinearEquationsSolution { particular, basis })
378 }
379
380 pub fn inverse(&self) -> Option<Matrix<R>> {
381 assert_eq!(self.shape.0, self.shape.1);
382 let n = self.shape.0;
383 if n >= 64 {
384 let m = n / 2;
385 let a = Self::new_with((m, m), |i, j| self[i][j].clone());
386 if let Some(mut ai) = a.inverse() {
387 let b = Self::new_with((m, n - m), |i, j| self[i][j + m].clone());
388 let c = Self::new_with((n - m, m), |i, j| self[i + m][j].clone());
389 let mut d = Self::new_with((n - m, n - m), |i, j| self[i + m][j + m].clone());
390 let u = &ai * &b;
391 let v = &c * &ai;
392 d -= &v * &b;
393 let di = d.inverse()?;
394 let r = &u * &di;
395 let t = &di * &v;
396 ai += &r * &v;
397 let mut inverse = Self::zeros((n, n));
398 for i in 0..m {
399 inverse[i][..m].clone_from_slice(&ai[i]);
400 for (x, y) in inverse[i][m..].iter_mut().zip(&r[i]) {
401 *x = R::neg(y);
402 }
403 }
404 for i in m..n {
405 for (x, y) in inverse[i][..m].iter_mut().zip(&t[i - m]) {
406 *x = R::neg(y);
407 }
408 inverse[i][m..].clone_from_slice(&di[i - m]);
409 }
410 return Some(inverse);
411 }
412 }
413 let mut a = self.clone();
414 let mut inverse = Self::eye((n, n));
415 let mut ranges: Vec<_> = (0..n).map(|i| (i, i + 1)).collect();
416 for r in 0..n {
417 let pivot = (r..n).find(|&i| !R::is_zero(&a[i][r]))?;
418 a.data.swap(r, pivot);
419 inverse.data.swap(r, pivot);
420 ranges.swap(r, pivot);
421
422 let d = R::inv(&a[r][r]);
423 for x in &mut a[r][r..] {
424 R::mul_assign(x, &d);
425 }
426 let (left, right) = ranges[r];
427 for x in &mut inverse[r][left..right] {
428 R::mul_assign(x, &d);
429 }
430
431 let (a_upper, a_lower) = a.data.split_at_mut(r + 1);
432 let pivot_a = &a_upper[r];
433 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r + 1);
434 let pivot_inverse = &inverse_upper[r];
435 let (ranges_upper, ranges_lower) = ranges.split_at_mut(r + 1);
436 let (left, right) = ranges_upper[r];
437 for ((a, inverse), range) in a_lower.iter_mut().zip(inverse_lower).zip(ranges_lower) {
438 if R::is_zero(&a[r]) {
439 continue;
440 }
441 let e = a[r].clone();
442 a[r] = R::zero();
443 R::add_scaled_assign(&mut a[(r + 1)..], &pivot_a[(r + 1)..], &R::neg(&e));
444 R::add_scaled_assign(
445 &mut inverse[left..right],
446 &pivot_inverse[left..right],
447 &R::neg(&e),
448 );
449 range.0 = range.0.min(left);
450 range.1 = range.1.max(right);
451 }
452 }
453 for r in (0..n).rev() {
454 let (left, right) = ranges[r];
455 let (inverse_upper, inverse_lower) = inverse.data.split_at_mut(r);
456 let pivot_inverse = &inverse_lower[0];
457 let (ranges_upper, _) = ranges.split_at_mut(r);
458 for ((a, inverse), range) in a.data[..r].iter_mut().zip(inverse_upper).zip(ranges_upper)
459 {
460 if R::is_zero(&a[r]) {
461 continue;
462 }
463 let e = a[r].clone();
464 a[r] = R::zero();
465 R::add_scaled_assign(
466 &mut inverse[left..right],
467 &pivot_inverse[left..right],
468 &R::neg(&e),
469 );
470 range.0 = range.0.min(left);
471 range.1 = range.1.max(right);
472 }
473 }
474 Some(inverse)
475 }
476
477 pub fn characteristic_polynomial(&mut self) -> Vec<R::T> {
478 let n = self.shape.0;
479 if n == 0 {
480 return vec![R::one()];
481 }
482 assert!(self.data.iter().all(|a| a.len() == n));
483 for j in 0..(n - 1) {
484 if let Some(x) = ((j + 1)..n).find(|&x| !R::is_zero(&self[x][j])) {
485 self.data.swap(j + 1, x);
486 self.data.iter_mut().for_each(|a| a.swap(j + 1, x));
487 let inv = R::inv(&self[j + 1][j]);
488 let mut v = vec![];
489 let src = std::mem::take(&mut self[j + 1]);
490 for a in self.data[(j + 2)..].iter_mut() {
491 let mul = R::mul(&a[j], &inv);
492 R::add_scaled_assign(&mut a[j..], &src[j..], &R::neg(&mul));
493 v.push(mul);
494 }
495 self[j + 1] = src;
496 for a in self.data.iter_mut() {
497 let v = R::dot_product(&a[(j + 2)..], &v);
498 R::add_assign(&mut a[j + 1], &v);
499 }
500 }
501 }
502 // dp[k][j - k] stores [x^k] det(xI - A[..j, ..j]).
503 let mut dp: Vec<Vec<R::T>> = (0..=n).map(|i| Vec::with_capacity(n + 1 - i)).collect();
504 dp[0].push(R::one());
505 for i in 0..n {
506 let mut c = vec![R::zero(); i + 1];
507 c[i] = R::neg(&self[i][i]);
508 let mut mul = R::one();
509 for j in (0..i).rev() {
510 mul = R::mul(&mul, &self[j + 1][j]);
511 c[j] = R::neg(&R::mul(&mul, &self[j][i]));
512 }
513 for k in (0..=i).rev() {
514 let mut value = R::dot_product(&dp[k], &c[k..]);
515 if k > 0 {
516 R::add_assign(&mut value, dp[k - 1].last().unwrap());
517 }
518 dp[k].push(value);
519 }
520 dp[i + 1].push(R::one());
521 }
522 dp.into_iter().map(|mut c| c.pop().unwrap()).collect()
523 }
524}
525
526impl<R> Index<usize> for Matrix<R>
527where
528 R: SemiRing,
529{
530 type Output = Vec<R::T>;
531 fn index(&self, index: usize) -> &Self::Output {
532 &self.data[index]
533 }
534}
535
536impl<R> IndexMut<usize> for Matrix<R>
537where
538 R: SemiRing,
539{
540 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
541 &mut self.data[index]
542 }
543}
544
545impl<R> Index<(usize, usize)> for Matrix<R>
546where
547 R: SemiRing,
548{
549 type Output = R::T;
550 fn index(&self, index: (usize, usize)) -> &Self::Output {
551 &self.data[index.0][index.1]
552 }
553}
554
555impl<R> IndexMut<(usize, usize)> for Matrix<R>
556where
557 R: SemiRing,
558{
559 fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
560 &mut self.data[index.0][index.1]
561 }
562}
563
564macro_rules! impl_matrix_pairwise_binop {
565 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident $(where [$($clauses:tt)*])?) => {
566 impl<R> $imp_assign for Matrix<R>
567 where
568 R: SemiRing,
569 $($($clauses)*)?
570 {
571 fn $method_assign(&mut self, rhs: Self) {
572 self.pairwise_assign(&rhs, |a, b| R::$method_assign(a, b));
573 }
574 }
575 impl<R> $imp_assign<&Matrix<R>> for Matrix<R>
576 where
577 R: SemiRing,
578 $($($clauses)*)?
579 {
580 fn $method_assign(&mut self, rhs: &Self) {
581 self.pairwise_assign(rhs, |a, b| R::$method_assign(a, b));
582 }
583 }
584 impl<R> $imp for Matrix<R>
585 where
586 R: SemiRing,
587 $($($clauses)*)?
588 {
589 type Output = Matrix<R>;
590 fn $method(mut self, rhs: Self) -> Self::Output {
591 self.$method_assign(rhs);
592 self
593 }
594 }
595 impl<R> $imp<&Matrix<R>> for Matrix<R>
596 where
597 R: SemiRing,
598 $($($clauses)*)?
599 {
600 type Output = Matrix<R>;
601 fn $method(mut self, rhs: &Self) -> Self::Output {
602 self.$method_assign(rhs);
603 self
604 }
605 }
606 impl<R> $imp<Matrix<R>> for &Matrix<R>
607 where
608 R: SemiRing,
609 $($($clauses)*)?
610 {
611 type Output = Matrix<R>;
612 fn $method(self, mut rhs: Matrix<R>) -> Self::Output {
613 rhs.pairwise_assign(self, |a, b| *a = R::$method(b, a));
614 rhs
615 }
616 }
617 impl<R> $imp<&Matrix<R>> for &Matrix<R>
618 where
619 R: SemiRing,
620 $($($clauses)*)?
621 {
622 type Output = Matrix<R>;
623 fn $method(self, rhs: &Matrix<R>) -> Self::Output {
624 let mut this = self.clone();
625 this.$method_assign(rhs);
626 this
627 }
628 }
629 };
630}
631
632impl_matrix_pairwise_binop!(Add, add, AddAssign, add_assign);
633impl_matrix_pairwise_binop!(Sub, sub, SubAssign, sub_assign where [R: SemiRing<Additive: Invertible>]);
634
635impl<R> Mul for Matrix<R>
636where
637 R: SemiRing,
638{
639 type Output = Matrix<R>;
640 fn mul(self, rhs: Self) -> Self::Output {
641 (&self).mul(&rhs)
642 }
643}
644impl<R> Mul<&Matrix<R>> for Matrix<R>
645where
646 R: SemiRing,
647{
648 type Output = Matrix<R>;
649 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
650 (&self).mul(rhs)
651 }
652}
653impl<R> Mul<Matrix<R>> for &Matrix<R>
654where
655 R: SemiRing,
656{
657 type Output = Matrix<R>;
658 fn mul(self, rhs: Matrix<R>) -> Self::Output {
659 self.mul(&rhs)
660 }
661}
662impl<R> Mul<&Matrix<R>> for &Matrix<R>
663where
664 R: SemiRing,
665{
666 type Output = Matrix<R>;
667 fn mul(self, rhs: &Matrix<R>) -> Self::Output {
668 assert_eq!(self.shape.1, rhs.shape.0);
669 if let Some(data) = R::try_matrix_product(&self.data, &rhs.data) {
670 return Matrix::from_vec(data);
671 }
672 let rhs = rhs.transpose();
673 Matrix::new_with((self.shape.0, rhs.shape.0), |i, j| {
674 R::dot_product(&self[i], &rhs[j])
675 })
676 }
677}
678
679fn strassen_rec<R: Ring>(
680 a: &[R::T],
681 b: &[R::T],
682 c: &mut [R::T],
683 shape: (usize, usize, usize),
684 stride_a: usize,
685 stride_b: usize,
686) {
687 let (n, m, p) = shape;
688 fn add_block<R: Ring>(
689 a: &[R::T],
690 b: &[R::T],
691 out: &mut [R::T],
692 n: usize,
693 stride_a: usize,
694 stride_b: usize,
695 ) {
696 for ((a, b), c) in a
697 .chunks(stride_a)
698 .zip(b.chunks(stride_b))
699 .zip(out.chunks_exact_mut(n))
700 {
701 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
702 *c = R::add(a, b);
703 }
704 }
705 }
706
707 fn sub_block<R: Ring>(
708 a: &[R::T],
709 b: &[R::T],
710 out: &mut [R::T],
711 n: usize,
712 stride_a: usize,
713 stride_b: usize,
714 ) {
715 for ((a, b), c) in a
716 .chunks(stride_a)
717 .zip(b.chunks(stride_b))
718 .zip(out.chunks_exact_mut(n))
719 {
720 for ((a, b), c) in a.iter().zip(b.iter()).zip(c.iter_mut()) {
721 *c = R::sub(a, b);
722 }
723 }
724 }
725
726 if n.min(m).min(p) <= 128 {
727 let transposed: Vec<_> = (0..p)
728 .flat_map(|j| (0..m).map(move |i| b[i * stride_b + j].clone()))
729 .collect();
730 for (a, c) in a.chunks(stride_a).zip(c.chunks_exact_mut(p)) {
731 for (b, c) in transposed.chunks_exact(m).zip(c) {
732 *c = R::dot_product(&a[..m], b);
733 }
734 }
735 return;
736 }
737 let (h, k, w) = (n / 2, m / 2, p / 2);
738 let a11 = 0;
739 let a12 = k;
740 let a21 = h * stride_a;
741 let a22 = a21 + k;
742 let b11 = 0;
743 let b12 = w;
744 let b21 = k * stride_b;
745 let b22 = b21 + w;
746
747 let block = h * w;
748 let mut buf = vec![R::zero(); h * k + k * w + block * 7];
749 let (s1, rest) = buf.split_at_mut(h * k);
750 let (s2, m_buf) = rest.split_at_mut(k * w);
751 let (m1, rest) = m_buf.split_at_mut(block);
752 let (m2, rest) = rest.split_at_mut(block);
753 let (m3, rest) = rest.split_at_mut(block);
754 let (m4, rest) = rest.split_at_mut(block);
755 let (m5, rest) = rest.split_at_mut(block);
756 let (m6, m7) = rest.split_at_mut(block);
757
758 // (A11 + A22)(B11 + B22)
759 add_block::<R>(&a[a11..], &a[a22..], s1, k, stride_a, stride_a);
760 add_block::<R>(&b[b11..], &b[b22..], s2, w, stride_b, stride_b);
761 strassen_rec::<R>(s1, s2, m1, (h, k, w), k, w);
762
763 // (A21 + A22) B11
764 add_block::<R>(&a[a21..], &a[a22..], s1, k, stride_a, stride_a);
765 strassen_rec::<R>(s1, &b[b11..], m2, (h, k, w), k, stride_b);
766
767 // A11 (B12 - B22)
768 sub_block::<R>(&b[b12..], &b[b22..], s2, w, stride_b, stride_b);
769 strassen_rec::<R>(&a[a11..], s2, m3, (h, k, w), stride_a, w);
770
771 // A22 (B21 - B11)
772 sub_block::<R>(&b[b21..], &b[b11..], s2, w, stride_b, stride_b);
773 strassen_rec::<R>(&a[a22..], s2, m4, (h, k, w), stride_a, w);
774
775 // (A11 + A12) B22
776 add_block::<R>(&a[a11..], &a[a12..], s1, k, stride_a, stride_a);
777 strassen_rec::<R>(s1, &b[b22..], m5, (h, k, w), k, stride_b);
778
779 // (A21 - A11)(B11 + B12)
780 sub_block::<R>(&a[a21..], &a[a11..], s1, k, stride_a, stride_a);
781 add_block::<R>(&b[b11..], &b[b12..], s2, w, stride_b, stride_b);
782 strassen_rec::<R>(s1, s2, m6, (h, k, w), k, w);
783
784 // (A12 - A22)(B21 + B22)
785 sub_block::<R>(&a[a12..], &a[a22..], s1, k, stride_a, stride_a);
786 add_block::<R>(&b[b21..], &b[b22..], s2, w, stride_b, stride_b);
787 strassen_rec::<R>(s1, s2, m7, (h, k, w), k, w);
788
789 let c11 = 0;
790 let c12 = w;
791 let c21 = h * p;
792 let c22 = c21 + w;
793 for ((((m1, m4), m5), m7), c) in m1
794 .iter()
795 .zip(m4.iter())
796 .zip(m5.iter())
797 .zip(m7.iter())
798 .zip(c[c11..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
799 {
800 *c = R::add(m1, m4);
801 R::sub_assign(c, m5);
802 R::add_assign(c, m7);
803 }
804 for ((m3, m5), c) in m3
805 .iter()
806 .zip(m5.iter())
807 .zip(c[c12..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
808 {
809 *c = R::add(m3, m5);
810 }
811 for ((m2, m4), c) in m2
812 .iter()
813 .zip(m4.iter())
814 .zip(c[c21..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
815 {
816 *c = R::add(m2, m4);
817 }
818 for ((((m1, m2), m3), m6), c) in m1
819 .iter()
820 .zip(m2.iter())
821 .zip(m3.iter())
822 .zip(m6.iter())
823 .zip(c[c22..].chunks_mut(p).flat_map(|c| c.iter_mut().take(w)))
824 {
825 *c = R::sub(m1, m2);
826 R::add_assign(c, m3);
827 R::add_assign(c, m6);
828 }
829}
830
831impl<R> Matrix<R>
832where
833 R: Ring,
834{
835 pub fn mul_strassen(&self, rhs: &Matrix<R>) -> Matrix<R> {
836 assert_eq!(self.shape.1, rhs.shape.0);
837 if let Some(data) = R::try_matrix_product(&self.data, &rhs.data) {
838 return Matrix::from_vec(data);
839 }
840 let (n, m) = self.shape;
841 let p = rhs.shape.1;
842 if n == 0 || m == 0 || p == 0 {
843 return Matrix::zeros((n, p));
844 }
845 let split = n.min(m).min(p).div_ceil(128).next_power_of_two();
846 if split <= 2 {
847 return self * rhs;
848 }
849 let rows = n.div_ceil(split) * split;
850 let inner = m.div_ceil(split) * split;
851 let cols = p.div_ceil(split) * split;
852 let mut a = vec![R::zero(); rows * inner];
853 for (a, data) in a.chunks_exact_mut(inner).zip(&self.data) {
854 a[..m].clone_from_slice(data);
855 }
856 let mut b = vec![R::zero(); inner * cols];
857 for (b, data) in b.chunks_exact_mut(cols).zip(&rhs.data) {
858 b[..p].clone_from_slice(data);
859 }
860 let mut c = vec![R::zero(); rows * cols];
861 strassen_rec::<R>(&a, &b, &mut c, (rows, inner, cols), inner, cols);
862 let mut res = Matrix::zeros((n, p));
863 for (data, c) in res.data.iter_mut().zip(c.chunks_exact(cols)) {
864 data.clone_from_slice(&c[..p]);
865 }
866 res
867 }
868}
869
870impl<R> MulAssign<&R::T> for Matrix<R>
871where
872 R: SemiRing,
873{
874 fn mul_assign(&mut self, rhs: &R::T) {
875 for i in 0..self.shape.0 {
876 for j in 0..self.shape.1 {
877 R::mul_assign(&mut self[(i, j)], rhs);
878 }
879 }
880 }More examples
crates/competitive/src/math/black_box_matrix.rs (line 175)
97 pub fn determinant(&self) -> R::T {
98 assert_eq!(self.shape.0, self.shape.1);
99 let n = self.shape.0;
100 let mut columns = vec![Vec::<(usize, R::T)>::new(); n];
101 for &(i, j, ref value) in &self.nonzero {
102 columns[j].push((i, value.clone()));
103 }
104 let mut degrees = vec![0; n];
105 for column in &mut columns {
106 column.sort_unstable_by_key(|&(i, _)| i);
107 let mut merged: Vec<(usize, R::T)> = Vec::with_capacity(column.len());
108 for (i, value) in column.drain(..) {
109 if let Some((_, x)) = merged.last_mut().filter(|(last, _)| *last == i) {
110 R::add_assign(x, &value);
111 } else {
112 merged.push((i, value));
113 }
114 }
115 merged.retain(|(i, value)| {
116 if R::is_zero(value) {
117 false
118 } else {
119 degrees[*i] += 1;
120 true
121 }
122 });
123 *column = merged;
124 }
125 let mut order: Vec<_> = (0..n).collect();
126 order.sort_unstable_by_key(|&j| columns[j].len());
127 let mut lower: Vec<Vec<(usize, R::T)>> = Vec::with_capacity(n);
128 let mut pivots: Vec<Option<usize>> = vec![None; n];
129 let mut x = vec![R::zero(); n];
130 let mut seen = vec![0; n];
131 let mut stack = Vec::new();
132 let mut support = Vec::new();
133 let mut determinant = R::one();
134 for (k, &j) in order.iter().enumerate() {
135 support.clear();
136 for &(i, _) in &columns[j] {
137 if seen[i] == k + 1 {
138 continue;
139 }
140 seen[i] = k + 1;
141 x[i] = R::zero();
142 stack.push((i, 0));
143 while let Some((i, next)) = stack.last_mut() {
144 if let Some(pivot) = pivots[*i]
145 && *next < lower[pivot].len()
146 {
147 let row = lower[pivot][*next].0;
148 *next += 1;
149 if seen[row] != k + 1 {
150 seen[row] = k + 1;
151 x[row] = R::zero();
152 stack.push((row, 0));
153 }
154 continue;
155 }
156 support.push(*i);
157 stack.pop();
158 }
159 }
160 for &(i, ref value) in &columns[j] {
161 x[i] = value.clone();
162 }
163 let mut pivot = None;
164 for &i in support.iter().rev() {
165 if let Some(p) = pivots[i] {
166 let factor = R::neg(&x[i]);
167 for &(row, ref value) in &lower[p] {
168 R::add_assign(&mut x[row], &R::mul(&factor, value));
169 }
170 } else if !R::is_zero(&x[i]) && pivot.is_none_or(|p| degrees[i] < degrees[p]) {
171 pivot = Some(i);
172 }
173 }
174 let Some(pivot) = pivot else { return R::zero() };
175 R::mul_assign(&mut determinant, &x[pivot]);
176 let inv = R::inv(&x[pivot]);
177 pivots[pivot] = Some(k);
178 lower.push(
179 support
180 .iter()
181 .filter(|&&i| pivots[i].is_none() && !R::is_zero(&x[i]))
182 .map(|&i| (i, R::mul(&x[i], &inv)))
183 .collect(),
184 );
185 }
186 for mut permutation in [order, pivots.into_iter().map(Option::unwrap).collect()] {
187 for i in 0..n {
188 while permutation[i] != i {
189 let j = permutation[i];
190 permutation.swap(i, j);
191 determinant = R::neg(&determinant);
192 }
193 }
194 }
195 determinant
196 }Dyn Compatibility§
This trait is not dyn compatible.
In older versions of Rust, dyn compatibility was called "object safety".