💰

Rustで動的計画法の実装：Coins

2023/04/14に公開

Rust

AtCoder

動的計画法

tech

はじめに

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

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

AからZまで問題が設定されているが、今回はIのCoins。 $N$ 枚のコインを投げて表の個数が裏の個数を上回る確率を求める比較的優しい問題。今回は、思わず貧乏性がでてメモリを節約してしまった話と、それが処理速度を落としてしまっていたという話になる。

利用したライブラリ

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

完成したコード

完成してから感じたが、動的計画法っぽくするのなら、 $i$ 個のコインを投げて $j$ 個表がでる確率をdp[i][j]として実装するのが正しい。ただ今回の問題の場合、確率なので復元も必要ないし、 $i$ 個のコインを投げた時の確率を求めるには、1つ前の $i-1$ 個のコインを投げた時の確率しか不要なので、1次元ベクトルdp[j]1本で実装してしまった。 $N\times N$ のテーブルを用意するのは勿体無いという貧乏性が抜けなくて困る。

coins.rs

use proconio::input;

fn main(){
    input!(
        n: usize,
        mut coins: [f64; n]
    );
    
    let target = (n / 2) + 1;
    let mut dp = vec![0.0; target + 1];

    dp[0] = 1.0;
    for i in 0..n {
        dp[target] = dp[target - 1] * coins[i] + dp[target];
        for j in (1 .. target).rev() {
            dp[j] = dp[j - 1] * coins[i] + dp[j] * (1.0 - coins[i]); 
        }
        dp[0] = dp[0] * (1.0 - coins[i]);
    }
    println!("{}", dp[target]);
}

アルゴリズム

アルゴリズムという程のものはないが、 $i-1$ 個のコインを投げて $j$ 個の表がでる確率が $dp[j]$ だった時、 $i$ 個目のコインが表なら、 $j+1$ 個目の表になるし、裏なら表の数は $j$ 個のままである。
このため、 $i$ 個目のコインを投げた時に $j$ 個のコインが表になっている確率は以下のようになる。

coins.rs

dp[j] = dp[j - 1] * coins[i] + dp[j] * (1.0 - coins[i]);

また、 $N$ 個のコインを投げ終わった後に「表の個数が多くなる」場合、つまり $\lfloor N/2 \rfloor + 1$ 個以上 (target個以上）の表が $i$ 個コインを投げている段階で既にでている時は、その後表だろうが裏だろうが最終的には表の個数が多くなる。このため、dp[target]をtarget枚以上の表がでている場合を表す事として、以下の式にする。

coins.rs

dp[target] = dp[target - 1] * coins[i] + dp[target];

降順でのループ

dp[j]を求めるために、dp[j]とdp[j-1]を利用している。今回の実装は $i$ 個コインを投げた時の確率で、 $i-1$ の時の確率を上書きしてしまうので、計算の順番を気をつけないといけない。具体的にはdp[j-1]はdp[j]の後で計算しないといけない。つまり、jは降順でループさせないといけない。

coins.rs

for j in (1 .. target).rev() {

Rustではrevを使って弱順にする。 1..5は 1, 2, 3, 4になるが、5..1は空になってしまう。このため。一度正順で作成したものをrev()で逆転させる必要がある。つまり(1..5).rev()で4, 3, 2, 1になる。5..1で期待するであろう5, 4, 3, 2を得るには(2..6).rev()になる。

2本のベクトルを使う

そもそも、1本のベクトルでやろうとしてたので順番を気にする必要があったのではないか。ならば、今（ $i-1$ 個のコインを投げた時）の確率列であるdp[]と、次（ $i$ 個のコインを投げた時）の確率列next_dp[]の2本を使った方が分かりやすい。

ループの最後にdpにnext_dpを入れているが、ここはベクトルのコピーではなく、ownerの変更なので重い処理では無い。

coins_2.rs

    dp[0] = 1.0;
    for i in 0..n {
        let mut next_dp = vec![0.0; n];
        next_dp[0] = dp[0] * (1.0 - coins[i]);
        for j in 1 .. target {
            next_dp[j] = dp[j - 1] * coins[i] + dp[j] * (1.0 - coins[i]); 
        }
        next_dp[target] = dp[target - 1] * coins[i] + dp[target];
        dp = next_dp;
    }

計算時間の測定

計算量としては変わらないはずであるが、2本のベクトルを使った方が毎回メモリ確保をしないといけないので遅いはず...と予測して計算量を測定した。問題を作るスクリプトは以下のような感じ。

gen_coins.py

import random
import sys

n = int(sys.argv[1])

p = [random.random() for i in range(n)]
print('{}'.format(n))
print(' '.join(map(lambda x: f'{x:.02f}', p)))

コインの数 $N$ を変化させて計算時間を測定した結果を以下に示す。予測と異なり2本のベクトルを使った方が計算時間は短い。

差がでてくる可能性としては以下の部分。ベクトルが1本の場合は、ループを順番に処理しないといけないが、2本のベクトルを使う場合は計算の順番に依存しないので並列化が可能。コンパイラがうまいこと並列化しなのではないか？と推測する。

coins.rs

for j in (1 .. target).rev() {
  dp[j] = dp[j - 1] * coins[i] + dp[j] * (1.0 - coins[i]); 
}

coins_2.rs

for j in 1 .. target {
  next_dp[j] = dp[j - 1] * coins[i] + dp[j] * (1.0 - coins[i]); 
}

上で説明文を書いていて感じたが、2本のベクトルを使ったコードの方が分かりやすい。つまり、結論としてはケチるとうまくいかないこともあるという話であった。

Rustで動的計画法の実装：Coins

はじめに

利用したライブラリ

完成したコード

アルゴリズム

降順でのループ

2本のベクトルを使う

計算時間の測定

関連記事

Discussion