cp_library_rs/linear_algrebra/
matrix_exp.rs1#![allow(clippy::needless_range_loop)]
4
5use num_traits::{One, Zero};
6
7#[derive(Debug, Clone)]
8pub struct Matrix<const N: usize, T>(pub [[T; N]; N]);
9
10impl<const N: usize, T> Matrix<N, T>
11where
12 T: One + Zero + Copy,
13{
14 pub fn id() -> Self {
16 let mut res = [[T::zero(); N]; N];
17 for i in 0..N {
18 res[i][i] = T::one();
19 }
20 Self(res)
21 }
22
23 pub fn apply(&self, v: &[T; N]) -> [T; N] {
25 let mut res = [T::zero(); N];
26 for i in 0..N {
27 for j in 0..N {
28 res[i] = res[i] + self.0[i][j] * v[j];
29 }
30 }
31 res
32 }
33
34 pub fn pow(&self, mut p: usize) -> Self {
36 let mut res = Self::id();
37 let mut tmp = self.clone();
38 while p > 0 {
39 if p & 1 == 1 {
40 res = tmp.dot(&res);
41 }
42 tmp = tmp.dot(&tmp);
43 p >>= 1;
44 }
45 res
46 }
47
48 pub fn dot(&self, other: &Self) -> Self {
50 let mut res = [[T::zero(); N]; N];
51 for i in 0..N {
52 for j in 0..N {
53 for k in 0..N {
54 res[i][j] = res[i][j] + self.0[i][k] * other.0[k][j];
55 }
56 }
57 }
58 Self(res)
59 }
60}