cp_library_rs/linear_algrebra/
matrix_exp.rs

1//! 行列累乗
2
3#![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    /// 単位行列を返す
15    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    /// ベクトル`x`と行列`A`について、`Ax`を返す
24    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    /// 行列の累乗を返す(繰り返し2乗法)
35    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    /// 行列のドット積
49    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}