🐸

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

2023/04/05に公開

Rust

AtCoder

動的計画法

tech

はじめに

Rustの勉強を行なっているが、実際にアルゴリズムを書くのがいいかな？と思い、基本である動的計画法を実装してみることにする。
丁度Educational DP Contestという動的計画法の練習を目的としたコンテストを見つけたのでやってみたいというのもある。

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

AからZまで問題が設定されているが、AのFrog1はBのFrog2を $k=2$ で固定した部分問題なのでFrog2から実装してみた。

利用したライブラリ

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

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

完成したコード

簡単な問題なのでいきなりコード。メモリを多めに確保してmの計算式を削除した方が速いのかもしれないけど、必要な分だけループを回すようにした。Visual Studio Codeで書いているが、型推論ができたところは補完されるのが面白い。

frog.rs

use proconio::input;

fn main(){
    input!(
        n: usize, k: usize,
        height: [i64; n],
    );
    let mut dp = vec![i64::MAX; n];
    dp[0] = 0;
    for i  in 0..(n - 1) {
        let m = min!(k + 1, n - i);
        for j in 1..m {
            chmin!(dp[i + j], dp[i] + (height[i] - height[i + j]).abs());
        }
    }
    println!("{}", dp[n-1]);
}

入力データの自動生成

コンテストのページに入力例と出力（正解）が載っているので、これで正しくインプリできているかが分かるが、折角なので、任意の問題を作ることができるスクリプトを作ってみる。

gen_input.py

#!/user/bin/env python
import random
import sys

n = int(sys.argv[1])
k = int(sys.argv[2])

h = [random.randint(1, 1000) for i in range(n)]
print('{} {}'.format(n, k))
print(' '.join(map(str, h)))

実行時間の計測

プログラムリストを見ると、計算オーダーは $O(n)$ なのは明確であるが、一応確認してみる。
シェルスクリプトを使って回してみる。

test.sh

#!/bin/bash

for k in 2 5 10; do
    for n in 10 20 50 100 200 500 1000 2000 5000 10000 20000 50000 100000; do
        python3 test_data.py $n $k | cargo r --release -q
    done
done

作業用メモリの初期化から答えの算出が終わるまでの時間を測定した。

dp_b.rs

    );
+   let start = time::Instant::now();
    let mut dp = vec![i64::MAX; n];
    
    }
-   println!("{}", dp[n-1]);
+   println!("{}, {}, {}", n, k, start.elapsed().as_nanos());
}

結果は以下の通り、完全ではないがだいたい $O(n)$ になっているのが分かる。 $k$ については $n$ が大きくなると無視できない感じ。ついでにi32とi64の比較も行った。64bit CPU上で動かしたのでi64の方が早いけど、 $n$ が小さい場合にいは差はほとんどなし。

計算時間の測定結果

読み込み時間

問題なのはどちらかと言えば、標準入力からの読み込み時間。最初のinput!にかかる時間を測定してみると、以下のような感じ。計算時間と比較しても1桁違う時間がかかっているし、単調増加ではない。I/O周りなので、試行回数が1回では不十分だと思うが、計算処理と比較してこっちがボトルネックになってしまうことは言える。

標準入力からの読込時間の測定結果