アルゴリズムと数学 057

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

これは状態遷移の問題ですね。タイルを左上から並べていきます。すなわち、最もタイルが左にある高さにタイルを敷き詰めます。こう決めることによって、タイルを並べる順番が一意に決まります。
状態を、各高さでタイルがどこまで並んでいるかを状態とします。ただし、一番小さい高さを0に正規化します。具体的にK=3として、最初は[0, 0, 0]で、次に縦と横にタイルを敷けるので、[2, 0, 0]と[1, 1, 0]が考えられます。[2, 0, 0]から縦に敷くと、[2, 1, 1]となりますが、これを正規化して[1, 0, 0]を状態とします。
だから、状態数はけっこう小さいです。なので、どの状態からどの状態へ遷移できるかを行列で表して、行列の(KN/2)乗を計算して、[0, 0, 0]から[0, 0, 0]へのパス数を計算します。より具体的には、[0, 0, 0]を0番として、0番に遷移できればいいので、べき乗した行列をPとすると、P[0][0]がパス数になります。

// Domino Tiling
#![allow(non_snake_case)]

use std::collections::{HashSet, HashMap};


//////////////////// constants ////////////////////

const D: i64 = 1000000007;


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

type Matrix = Vec<Vec<i64>>;

fn mat_mul(A: &Matrix, B: &Matrix, d: i64) -> Matrix {
    let L = A.len();
    let M = B.len();
    let N = B[0].len();
    (0..L).map(|i| (0..N).
            map(|j| (0..M).map(|k| (A[i][k] * B[k][j]) % d).sum::<i64>() % d).
                                                    collect::<Vec<i64>>()).
                                                    collect::<Matrix>()
}

fn mat_pow(A: &Matrix, e: u64, d: i64) -> Matrix {
    if e == 1 {
        A.clone()
    }
    else if e % 2 == 0 {
        let B = mat_pow(&A, e/2, d);
        mat_mul(&B, &B, d)
    }
    else {
        mat_mul(&A, &mat_pow(&A, e-1, d), d)
    }
}


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

type State = Vec<i64>;

fn read_input() -> (usize, usize) {
    let v = read_vec();
    (v[0], v[1])
}

fn normalize(s0: State) -> State {
    let min_x = s0.iter().min().unwrap();
    s0.iter().map(|x| x - min_x).collect::<State>()
}

fn add(s0: &State, y: usize, x: i64) -> State {
    let mut s = s0.clone();
    if x == 1 {
        s[y] += 1;
        s[y+1] += 1
    }
    else {
        s[y] += 2
    }
    normalize(s)
}

fn nexts(s0: &State) -> Vec<State> {
    let H = s0.len();
    let mut states: Vec<State> = Vec::new();
    let y = s0.iter().enumerate().filter(|(_i, &x)| x == 0).
                                    map(|(i, &_x)| i).next().unwrap();
    states.push(add(s0, y, 2));
    if y < H - 1 && s0[y+1] == 0 {
        states.push(add(s0, y, 1))
    }
    states
}

fn state_to_i64(s: &State) -> i64 {
    s.iter().fold(0, |x, y| x * 3 + y)
}

fn collect_edges(K: usize) -> Vec<(i64, i64)> {
    let mut edges: Vec<(i64, i64)> = Vec::new();
    let s0 = vec![0; K];
    let mut stack = vec![s0];
    let mut visited: HashSet<i64> = HashSet::new();
    visited.insert(0);
    while !stack.is_empty() {
        let s = stack.pop().unwrap();
        let neighbors = nexts(&s);
        let hash = state_to_i64(&s);
        for neighbor in neighbors.into_iter() {
            let hash1 = state_to_i64(&neighbor);
            if !visited.contains(&hash1) {
                visited.insert(hash1.clone());
                stack.push(neighbor)
            }
            edges.push((hash, hash1));
        }
    }
    edges
}

fn create_dict(edges: &Vec<(i64, i64)>) -> HashMap<i64, usize> {
    let mut dic: HashMap<i64, usize> = HashMap::new();
    for &(v, w) in edges.iter() {
        if !dic.contains_key(&v) {
            dic.insert(v, dic.len());
        }
        if !dic.contains_key(&w) {
            dic.insert(w, dic.len());
        }
    }
    dic
}

fn make_matrix_core(edges: &Vec<(i64, i64)>, dic: &HashMap<i64, usize>)
                                                            -> Matrix {
    let L = dic.len();
    let mut A: Matrix = (0..L).map(|_| vec![0; L]).collect();
    for &(v, w) in edges.iter() {
        let i = dic[&v];
        let j = dic[&w];
        A[i][j] = 1;
    }
    A
}

fn make_matrix(K: usize) -> Matrix {
    let edges = collect_edges(K);
    let dic = create_dict(&edges);
    make_matrix_core(&edges, &dic)
}

fn f(K: usize, N: usize) -> i64 {
    let A = make_matrix(K);
    let P = mat_pow(&A, (K*N/2) as u64, D);
    P[0][0]
}

fn main() {
    let (K, N) = read_input();
    println!("{}", f(K, N))
}