Rustで動的計画法 (2)

『アルゴリズムとデータ構造』を Rust で実装してみるシリーズ。5章の動的計画法。

前回 → https://quantized-cube.com/rust-dp-01/

Frog問題を一気に実装する。

use proconio::input;

const INF: i64 = 1 << 60; // 十分大きい値とする (ここでは 2^60)

fn input(use_default_values: bool) -> (usize, Vec<i64>) {
    if use_default_values {
        let n: usize = 7;
        let h: Vec<i64> = vec![2, 9, 4, 5, 1, 6, 10];
        return (n, h);
    } else {
        println!("Input n and h:");
        input! {
            n: usize,
            h: [i64; n],
        }
        return (n, h);
    }
}

fn chmin<T: PartialOrd>(a: &mut T, b: T) {
    if *a > b {
        *a = b;
    }
}

/// 「緩和」を意識した動的計画法
pub fn code_5_3(use_default_values: bool) -> i64 {
    // 入力
    let (n, h) = input(use_default_values);

    // 初期化 (最小化問題なので INF に初期化)
    let mut dp = vec![INF; n];

    // 初期条件
    dp[0] = 0;

    // ループ
    for i in 1..n {
        let (left, right) = dp.split_at_mut(i);
        let a = &mut right[0];
        let b = left[i - 1] + (h[i] - h[i - 1]).abs();
        chmin(a, b);
        if i > 1 {
            let b = left[i - 2] + (h[i] - h[i - 2]).abs();
            chmin(a, b);
        }
    }

    // 答え
    dp[n - 1]
}

/// 「配る遷移形式」
pub fn code_5_4(use_default_values: bool) -> i64 {
    // 入力
    let (n, h) = input(use_default_values);

    // 初期化 (最小化問題なので INF に初期化)
    let mut dp = vec![INF; n];

    // 初期条件
    dp[0] = 0;

    // ループ
    for i in 0..n {
        if i + 1 < n {
            let b = dp[i] + (h[i] - h[i + 1]).abs();
            chmin(&mut dp[i + 1], b);
        }
        if i + 2 < n {
            let b = dp[i] + (h[i] - h[i + 2]).abs();
            chmin(&mut dp[i + 2], b);
        }
    }

    // 答え
    dp[n - 1]
}

/// 再帰関数を用いる単純な全探索
fn rec_5_5(i: usize, h: &Vec<i64>) -> i64 {
    // 足場 0 のコストは 0
    if i == 0 {
        return 0;
    }

    // 答えを格納する変数を INF に初期化する
    let mut res = INF;

    // 頂点 i - 1 から来た場合
    let b = (rec_5_5(i - 1, h) + h[i] - h[i - 1]).abs();
    chmin(&mut res, b);

    // 頂点 i - 2 から来た場合
    if i > 1 {
        let b = (rec_5_5(i - 2, h) + h[i] - h[i - 2]).abs();
        chmin(&mut res, b);
    }

    // 答えを返す
    return res;
}

pub fn code_5_5(use_default_values: bool) -> i64 {
    // 入力
    let (n, h) = input(use_default_values);

    rec_5_5(n - 1, &h)
}

/// 「メモ化再帰」で解く
fn rec_5_6(i: usize, h: &Vec<i64>, dp: &mut Vec<i64>) -> i64 {
    // DP の値が更新されていたらそのままリターン
    if dp[i] < INF {
        return dp[i];
    }

    // ベースケース: 足場 0 のコストは 0
    if i == 0 {
        return 0;
    }

    // 答えを表す変数を INF で初期化する
    let mut res = INF;

    // 足場 i - 1 から来た場合
    let b = rec_5_6(i - 1, h, dp) + (h[i] - h[i - 1]).abs();
    chmin(&mut res, b);

    // 足場 i - 2 から来た場合
    if i > 1 {
        let b = rec_5_6(i - 2, h, dp) + (h[i] - h[i - 2]).abs();
        chmin(&mut res, b);
    }

    // 結果をメモしながら、返す
    dp[i] = res;
    res
}

pub fn code_5_6(use_default_values: bool) -> i64 {
    // 入力
    let (n, h) = input(use_default_values);

    // 初期化 (最小化問題なので INF に初期化)
    let mut dp = vec![INF; n];

    // 答え
    rec_5_6(n - 1, &h, &mut dp)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn code_5_3_works() {
        let result = code_5_3(true);
        assert_eq!(result, 8);
    }

    #[test]
    fn code_5_4_works() {
        let result = code_5_4(true);
        assert_eq!(result, 8);
    }

    #[test]
    fn code_5_5_works() {
        let result = code_5_5(true);
        assert_eq!(result, 8);
    }

    #[test]
    fn code_5_6_works() {
        let result = code_5_6(true);
        assert_eq!(result, 8);
    }
}

split_at_mut を使わずにいい感じに書く方法はないだろうか？

GitHub のリポジトリはこちら https://github.com/quantized-cube/book_algorithm_solution_rust