🍣

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

2023/04/16に公開

Rust

AtCoder

動的計画法

tech

はじめに

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

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

AからZまで問題が設定されているが、今回はJのSushi。問題としては難易度が高めだが、実装自体はそんなに難しくない。

利用したライブラリ

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

完成したコード

実装としては、キャッシュ付きの再帰法である。グローバル変数を避けるために構造体を作ってキャッシュはそこで保持するようにしている。1皿に乗る寿司の最大数も可変にしようかと思ったが、dpの次元を変えないといけないのでやめた。

sushi.rs

use proconio::input;

const SUSHI_MAX : usize = 3;

struct Solver {
    dp: Vec<Vec<Vec<Option<f64>>>>,
    n: f64
}

impl Solver {
    fn new(n:usize) -> Self {
        let mut obj = Solver { dp: vec![vec![vec![None; n + 1]; n + 1]; n + 1], n:n as f64};
        obj.dp[0][0][0] = Some(0.0);
        obj
    }

    fn solve(&mut self, i:usize, j:usize, k:usize) -> f64 {
        if let Some(value) = self.dp[i][j][k] { return value; };
        let denominator = (i + j + k) as f64;
        let mut result = self.n;
        if i > 0 { result += i as f64 * self.solve(i - 1, j + 1, k); };
        if j > 0 { result += j as f64 * self.solve(i, j - 1, k + 1); };
        if k > 0 { result += k as f64 * self.solve(i, j, k - 1); };
        result = result / denominator;
        self.dp[i][j][k] = Some(result);
        return result;
    }
}

fn main(){
    input! {
        n: usize,
        dish: [usize; n]
    }

    let mut sushi = vec![0; SUSHI_MAX + 1];
    for d in &dish {
        assert!(*d <= SUSHI_MAX);
        sushi[*d] += 1;
    }
    let mut solver = Solver::new(n);
    let _result = solver.solve(sushi[3], sushi[2], sushi[1]);
    println!("{}", result);
}

アルゴリズム

これは漸化式をたてるまでが難しい。別の問題であるがガチャのコンプリート問題で基本的な考え方を説明する。全部同じ確率 $1/n$ で排出される $n$ 種類のアイテムがあるガチャがあるとして、コンプリートするまでに、何回ガチャを回す必要があるか？という問題がある。

現在 $m$ 個のアイテムを持っている時、あと何回ガチャを回すとコンプリートできるか？の期待値を $E[X_m]$ とする。 $m$ 個のアイテムを持った状態で、1回ガチャを回して、持っていないものがでる確率は $(n-m)/n$ で、その後ガチャを回す回数の期待値は $E[X_{m+1}]$ 。持っているものがでる確率は $m/n$ で、その後ガチャを回す回数の期待値は変化しないので $E[X_m]$ となる。つまり、以下の漸化式ができる。

E[X_m] = 1 + \left( \frac{n-m}{n}E[X_{m+1}] + \frac{m}{n}E[X_{m}]\right)

両編に $E[X_m]$ があるので、変形して漸化式を作り、何も持っていない時からの期待値 $E[X_0]$ を求める。

この問題もこれをベースに、dp[i][j][k]を3貫, 2貫, 1貫寿司がのった皿の枚数がi, j, kの時に、あと何回で全部の寿司を食べることができるか？の期待値だとして漸化式を立てれば良い。

キャスト（型変換）

Rustでは型変換は明示的に行う必要がある上に書き方はちょっと独特。swiftもこんな書き方だったので最近の言語での流行りなのかもしれない。難しい事はないのだが、 $i, j, k$ はベクトルの添字usizeとなるので、期待値の計算を行う場合は浮動小数点f64に変換しておく必要がある。

sushi.rs

let denominator = (i + j + k) as f64;

Optionを使ったキャッシュ

同じ計算を2度実施しなくていいように、キャッシュに保存しておく。そもそもキャッシュに保存されているかどうかを明確にするために、キャッシュを保存していない事を表現するための-1や極端に大きな値を初期値として与える場合が多い。今回は汎用的な方法として、Optionを使ってみた。

初期値としては値がない事を示すNoneを与え、キャッシュ値としてSome(T)を使う。「特別な値」を使わなくてすむし、見やすいのでいいかと思う。

sushi.rs

fn new(n:usize) -> Self {
        let mut obj = Solver { dp: vec![vec![vec![None; n + 1]; n + 1]; n + 1], n:n as f64};
        obj.dp[0][0][0] = Some(0.0);
        obj
    }

async / await

今回、これを非同期通信で書くとどうなるのか？非同期通信回りの実装を学習してみたいし、総当たりなのでマルチスレッドにすると速くなるのではと思って、async/awaitで動作するように改修してみた。

排他制御

総当たりといいつつ、キャッシュを使って同じ計算をする必要がないようにしているので、キャッシュは、すべてのスレッドからアクセスできる必要がある。つまり、排他制御をきちんと実施する必要がある。複数スレッドからのアクセスを実現するために参照を複数作成できるArc<>と、読み書きのロックがかけられるRwLock<>の組み合わせがよく使われる。Mutexでもいいのだが、読み書き別々にロックされるRwLock<>の方を使ってみた。

作り方は簡単で、ターゲットのオブジェクトを入れ子にして生成する。

sushi_async.rs

let cache = Arc::new(RwLock::new(Cache::new(n)));

Arcは、clone()することで参照を複数作成できる。

sushi_async.rs

solve(cache.clone(), sushi[3], sushi[2], sushi[1], true).await

使う時はread()かwrite()で所有権を取得してからアクセスすれば良い。

sushi_async.rs

let cache_read = cache.read();
if let Some(value) = cache_read.dp[i][j][k] { return value; };

再帰関数でのasync/await

再帰関数でasync/awaitを使うためには、async_recursionを使うといい。本当はすごく複雑な記述をする必要があるのだが、#[async_recursion]を置いておけば、コンパイラがやってくれる。

sushi_async.rs

#[async_recursion]
async fn solve(cache:Arc<RwLock<Cache>>, i:usize, j:usize, k:usize, is_span:bool) -> f64 {

計算時間の測定

async/awaitを使った場合と、使わなかった場合での計算時間を測定した。async/awaitを使った方が断トツに遅い。

tokio-consoleで見てみたが、再帰呼び出し毎にtaskを作る実装だと121,389もtaskができていて、taskのスケジューリング処理が原因かと思う。最初のところだけスレッドにするなど、taskを減らすと変わるかもしれない。

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

はじめに

利用したライブラリ

完成したコード

アルゴリズム

キャスト（型変換）

Optionを使ったキャッシュ

async / await

排他制御

再帰関数でのasync/await

計算時間の測定

関連記事

Discussion