cp_library_rs/linear_algrebra/
dynamic_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<T> {
9    pub R: usize,
10    pub C: usize,
11    pub arr: Vec<Vec<T>>,
12}
13
14impl<T> Matrix<T>
15where
16    T: One + Zero + Copy,
17{
18    pub fn new(data: Vec<Vec<T>>) -> Self {
19        let R = data.len();
20        let C = data[0].len();
21        Self { R, C, arr: data }
22    }
23
24    /// 単位行列を返す
25    pub fn id(&self) -> Self {
26        let mut res = vec![vec![T::zero(); self.C]; self.R];
27        for i in 0..self.R {
28            res[i][i] = T::one();
29        }
30        Self {
31            R: self.R,
32            C: self.C,
33            arr: res,
34        }
35    }
36
37    /// ベクトル`x`と行列`A`について、`Ax`を返す
38    pub fn apply(&self, v: &[T]) -> Vec<T> {
39        assert_eq!(v.len(), self.C);
40        let mut res = vec![T::zero(); self.R];
41        for i in 0..self.R {
42            for j in 0..self.R {
43                res[i] = res[i] + self.arr[i][j] * v[j];
44            }
45        }
46        res
47    }
48
49    /// 行列の累乗を返す（繰り返し2乗法）
50    pub fn pow(&self, mut p: usize) -> Self {
51        let mut res = self.id();
52        let mut tmp = self.clone();
53        while p > 0 {
54            if p & 1 == 1 {
55                res = tmp.dot(&res);
56            }
57            tmp = tmp.dot(&tmp);
58            p >>= 1;
59        }
60        res
61    }
62
63    /// 行列`A`,`B`のドット積`AB`を求める
64    pub fn dot(&self, other: &Self) -> Self {
65        assert_eq!(self.C, other.R);
66        let mut res = vec![vec![T::zero(); other.C]; self.R];
67        for i in 0..self.R {
68            for j in 0..other.C {
69                for k in 0..self.C {
70                    res[i][j] = res[i][j] + self.arr[i][k] * other.arr[k][j];
71                }
72            }
73        }
74        Self {
75            R: self.R,
76            C: other.C,
77            arr: res,
78        }
79    }
80}
cp_library_rs/linear_algrebra/dynamic_matrix_exp.rs

cp_library_rs/linear_algrebra/
dynamic_matrix_exp.rs
