AtCoder Beginner Contest 260 C

https://atcoder.jp/contests/abc260/tasks/abc260_c

結局、HashMapはentryを使わないとうまくいかないですね。

// Changing Jewels
#![allow(non_snake_case)]

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 ////////////////////

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

fn mat_mul(A: &Matrix<usize>, B: &Matrix<usize>) -> Matrix<usize> {
    let H = A.len();
    let L = B.len();
    let W = B[0].len();
    let C: Matrix<usize> = (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(A: &Matrix<usize>, e: usize) -> Matrix<usize> {
    if e == 0 {
        let N = A.len();
        (0..N).map(|i|
            (0..N).map(|j| if i == j { 1 } else { 0 }).collect()).collect()
    }
    else if e == 1 {
        A.to_vec()
    }
    else if e % 2 == 0 {
        let B = mat_pow(A, e/2);
        mat_mul(&B, &B)
    }
    else {
        mat_mul(&A, &mat_pow(A, e-1))
    }
}

fn mat_print(A: &Matrix<usize>) {
    for v in A.iter() {
        println!("{:?}", v)
    }
}

//////////////////// main ////////////////////

fn main() {
    let v = read_vec();
    let N: usize = v[0];
    let X: usize = v[1];
    let Y: usize = v[2];
    let A: Matrix<usize> = vec![ vec![X+1, 1], vec![X*Y, Y] ];
    let B = mat_pow(&A, N - 1);
    println!("{}", B[1][0])
}