アルゴリズムと数学 100

https://atcoder.jp/contests/math-and-algorithm/tasks/math_and_algorithm_bv

 M = \begin{pmatrix}
1-X & Y & 0 \\
0 & 1-Y & Z \\
X & 0 & 1-Z \\
\end{pmatrix}
という行列を考えると、1秒進めると
 \begin{pmatrix}
a \\
b \\
c \\
\end{pmatrix} \leftarrow M \begin{pmatrix}
a \\
b \\
c \\
\end{pmatrix}
となります。T秒後には、
 \begin{pmatrix}
a \\
b \\
c \\
\end{pmatrix} \leftarrow M^T \begin{pmatrix}
a \\
b \\
c \\
\end{pmatrix}
となります。
行列の乗算、べき算のコードをいろんな数値で成り立つようにしてみました。AIが書いてくれるから簡単ですね。

// Simulation of Chemicals
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// Matrix ////////////////////

use std::ops::{Mul, Rem, Sub, Div};
use std::cmp::PartialEq;
use std::iter::Sum;

type Matrix<T> = Vec<Vec<T>>;

fn mat_mul<T: Copy + Mul<Output = T> + Sum>(A: &Matrix<T>,
                                            B: &Matrix<T>) -> Matrix<T> {
    let H = A.len();
    let L = B.len();
    let W = B[0].len();
    let C: Matrix<T> = (0..H).map(|i|
                       (0..W).map(|j|
                       (0..L).map(|k| A[i][k] * B[k][j]).sum()
                       ).collect()
                       ).collect();
    return C
}

fn mat_pow<T: Copy + Mul<Output = T> + Sum + From<u8>,
           E: Copy + Rem<Output = E> + Sub<Output = E> + Div<Output = E> +
                PartialEq + From<u8>>
        (A: &Matrix<T>, e: E) -> Matrix<T> {
    let zero = E::from(0);
    let one = E::from(1);
    let two = E::from(2);
    
    if e == zero {
        let N = A.len();
        (0..N).map(|i|
        (0..N).map(|j| if i == j { T::from(1) } else { T::from(0) }
        ).collect()
        ).collect()
    }
    else if e == one {
        A.to_vec()
    }
    else if e % two == zero {
        let B = mat_pow::<T, E>(A, e / two);
        mat_mul(&B, &B)
    }
    else {
        mat_mul(&A, &mat_pow::<T, E>(A, e - one))
    }
}


//////////////////// process ////////////////////

type Query = (f64, f64, f64, i32);

fn read_query() -> Query {
    let v = read_vec();
    let (X, Y, Z) = (v[0], v[1], v[2]);
    let T = v[3] as i32;
    (X, Y, Z, T)
}

fn make_matrix(X: f64, Y: f64, Z: f64) -> Matrix<f64> {
    let M: Matrix<f64> = vec![vec![1.0-X, Y, 0.0],
                              vec![0.0, 1.0-Y, Z],
                              vec![X, 0.0, 1.0-Z]];
    M
}

fn f_each(query: Query) {
    let (X, Y, Z, T) = query;
    let M = make_matrix(X, Y, Z);
    let M_all = mat_pow(&M, T);
    let v: Matrix<f64> = vec![vec![1.0]; 3];
    let w = mat_mul(&M_all, &v);
    let (A, B, C) = (w[0][0], w[1][0], w[2][0]);
    println!("{} {} {}", A, B, C)
}

fn f(Q: i32) {
    for _ in 0..Q {
        let query = read_query();
        f_each(query)
    }
}

fn main() {
    let Q: i32 = read();
    f(Q)
}