🫥

Rustで動的計画法：Slimes

2023/04/25に公開

Rust

AtCoder

動的計画法

tech

はじめに

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

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

AからZまで問題が設定されているが、今回はNのSlimesを解いてみる。

利用したライブラリ

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

久々にchmin!も使う。

完成したコード

dequeueの時と同じように、dp[i][j]を $j$ から $i$ までのスライムを合体する時に必要なコストの最小値として解いている。今回も下三角行列しか使わないので、メモリは積極的にケチっている。

slimes.rs

use proconio::input;

fn main(){
    input!(
        n: usize,
        slimes: [u64; n]
    );

    let mut dp = Vec::with_capacity(n);
    for i in 0..n {
        dp.push(vec![0u64; i+1]);
    }
    let cost = |j: usize, i: usize| -> u64 { (&slimes[j..i+1]).iter().sum() };
 
    for j in 1..n {
        for i in j..n {
            let j_x = i - j; // (dp[j].len() - 1) - i
            dp[i][j_x] = dp[j_x][j_x] + dp[i][j_x + 1];
            for k in 1..j {
                chmin!(dp[i][j_x], dp[j_x + k][j_x] + dp[i][j_x + k + 1]);
            }
            dp[i][j_x] += cost(j_x, i);
        }
    }
    println!("{}, {}, {}", n, start.elapsed().as_nanos(), dp[n - 1][0]);
}

アルゴリズム

$j$ から $i$ までのスライムの1つ前の状態は、 $j$ から $k$ までと $k+1$ から $i$ （ $j\leq k < i$ )までのスライムだったという事である。つまり、dp[i][j]を求めるには、d[i][k+1] + d[k][j]のコストが最も小さい組み合わせを見つければ良い。

dp[i][j]を計算するまでに、必要なコストを全て計算するためには、下三角行列の斜面、行列の対角線から順番に斜めに計算していけば良い。ループでは書きにくいのでj_xを作って、これを $j$ の代わりに使っている。

なお、dp[]はゼロで初期化しているので最初の1つ (k=0) だけはベタに計算して代入している。

イテレータの活用

$j$ から $i$ までのスライムを合体してできるスライムの大きさは、一意に決まるので、costという関数 (closure) で求めるようにしている。ここの計算結果を保存しながら進めるのもいいのだが、 $N$ がそこまで大きくなく、大した計算量ではないので逐次計算するようにした。

ここはforで実装するとすると以下のようになると思うが、iter()を使った方がすっきり書ける。おそらく処理時間はそんなに変わらないと思うが。

loopでの実装

let cost = |i: usize, j: usize| -> u64 { 
    let mut sum = 0;
    for k in i..j+1 {
        sum += slimes[k];
    }
    sum
}

計算時間の測定

$N$ を50から400まで変化させた時の計算処理時間を測定した結果を以下に示す。3重ループなので $O(N^3)$ であるが、 $N$ が400までなので、数十ms程度で何ら問題はない。 $N$ がもっと大きくなった場合はどうするか？については、今の所いい案が思い浮かばないのが困ったものだが...