Rustで動的計画法 (3)

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

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

今回はナップサック問題。

use proconio::input;

fn input(use_default_values: bool) -> (usize, usize, Vec<usize>, Vec<i64>) {
    if use_default_values {
        let n: usize = 6;
        let w: usize = 7;
        let weight: Vec<usize> = vec![2, 1, 3, 2, 1, 5];
        let value: Vec<i64> = vec![3, 2, 6, 1, 3, 85];
        return (n, w, weight, value);
    } else {
        println!("Input n, w, and weight_value_pair:");
        input! {
            n: usize,
            w: usize,
            weight_value_pair: [(usize, i64); n],
        }
        let mut weight: Vec<usize> = Vec::new();
        let mut value: Vec<i64> = Vec::new();
        for item in weight_value_pair.into_iter() {
            weight.push(item.0);
            value.push(item.1);
        }
        return (n, w, weight, value);
    }
}

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

/// ナップサック問題に対する動的計画法
pub fn code_5_7(use_default_values: bool) -> i64 {
    // 入力
    let (n, w, weight, value) = input(use_default_values);

    // DP テーブル定義
    let mut dp = vec![vec![0; w + 1]; n + 1];

    // DPループ
    for i in 0..n {
        for u in 0..=w {
            // i 番目の品物を選ぶ場合
            if u >= weight[i] {
                let b = dp[i][u - weight[i]] + value[i];
                chmax(&mut dp[i + 1][u], b);
            }
            // i 番目の品物を選ばない場合
            let b = dp[i][u];
            chmax(&mut dp[i + 1][u], b);
        }
    }

    // 最適値の出力
    dp[n][w]
}

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

    #[test]
    fn code_5_7_works() {
        let result = code_5_7(true);
        assert_eq!(result, 90);
    }
}

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