🌴

Rustで動的計画法の実装 : Vacation

2023/04/06に公開

Rust

AtCoder

動的計画法

tech

はじめに

動的計画法を実装してみて、Rustの勉強をやってみる。問題としてはEducational DP Contestという動的計画法の練習を目的としたコンテストのものを使用。

https://atcoder.jp/contests/dp/

AからZまで問題が設定されているが、今回はCのVacation。1日3種類の行動の中から選択して休暇中の幸福度を最大化する問題。ただし、二日連続で同じ行動はできない。

利用したライブラリ

chmaxは検索するといくつかでてくるので、それを利用。

標準入力からデータを読み取る部分はproconioを利用。

完成したコード

「3種類の行動」ということになっていたけど、そこも変更可能にしたいよね。ということでACTION_NUMを変更することで行動が３つ以外の問題にも対応した。

vacation.rs

use itertools::Itertools;

const ACTION_NUM : usize = 3;
fn main(){
    input!(
        n: usize,
        happiness: [[i64; ACTION_NUM]; n],
    );

    let mut dp = vec![vec![0; ACTION_NUM]; n];
    dp[0] = happiness[0].clone();
    for i in 1..n  {
        for x in (0..ACTION_NUM).permutations(2) {
            chmax!(dp[i][x[0]], dp[i-1][x[1]] + happiness[i][x[0]]);
        }
    }
    let answer = dp[n-1].iter().max().unwrap();
    println!("{}", answer);
}

基本的なアルゴリズム

$i$ 日目に行動 $j$ を行った時の幸福度の最大値をdp[i][j]とする^[1]。 $i$ 日目に行動 $j$ を行った時の幸福度はhappiness[i][j]に入っているので、以下のように書ける。二日連続で同じ行動はできないので、前日の行動 $k$ と今日の行動 $j$ は違う(i != k)必要がある。

vacation_original.rs

        for j in 0..ACTION_NUM {
            for k in 0..ACTION_NUM {
                if j != k {
                    chmax!(dp[i][j], dp[i-1][k] + happiness[i][j]);
                }    
            }
        }

ただ、2重ループになるし、条件文も入るのでちょっとソースの見通しが悪い。そこで最終的には以下のようにした。

vacation.rs

        for x in (0..ACTION_NUM).permutations(2) {
            chmax!(dp[i][x[0]], dp[i-1][x[1]] + happiness[i][x[0]]);
        }

permutationsはitertools::Itertoolsに含まれており、順列を返す。つまり ${}_n P_k$ 。基本的にfor文と一緒に以下のように使う。

permutations

    for e in (0..3).permutations(2) {
        print!("{:?}", e);
    }

output

[0, 1][0, 2][1, 0][1, 2][2, 0][2, 1]

二日連続で同じ行動をすることを許容するなら直積(product)を使えば良いが、そもそもその時は、毎日一番幸福度の高い行動をとる貪欲な方法で最適解が出る。

結果

結果は $n$ 日目(dp[n-1])で一番高い幸福度を求める。

vacations.rs

    let answer = dp[n-1].iter().max().unwrap();

max()の返り値としては、エラーかどうかも返すためにOption<&T>型になっている。正常時であればSome(&17)のように帰ってくるし、空行列だったりエラーがあった場合は、Noneが帰ってくる。Some(&17)を&17にするためにunwrap()が必要になる。

初期値

1日目はどの行動を行なってもいいので、happiness[0]をそのままコピーする。

vacations.rs

    let mut dp = vec![vec![0; ACTION_NUM]; n];
    dp[0] = happiness[0].clone();

最初は、vecのvecではなく、配列のvecにしていたのではまってしまった。

    let mut dp = vec![[0; ACTION_NUM]; n];
    dp[0] = happiness[0];

dp[0]は配列だけど、happiness[0]はvec型なので入力ができずにコンパイルできなかった。何か方法があるのかもしれないが、分からなかったので、vecのvecとした。

dp[0]は0で初期化して、dp[1]を1日目として動かしても動作するし、問題文には合っている気がする。しかし、happinessが0-originになっているので、それに合わせている。0-originと1-originを混合させるのはバグの元になると思う。

脚注

正確には0-originなのでdp[i-1][j] ↩︎