🍬

Rustで動的計画法：Candies

2023/04/22に公開

Rust

AtCoder

動的計画法

tech

はじめに

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

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

AからZまで問題が設定されているが、今回はMのCandiesを解いてみる。 $K$ 個のキャンディを $N$ 人の子供に配る問題。そんなに難しく無いように見えるが

利用したライブラリ

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

後半では並列化のためにrayonを利用。
https://docs.rs/rayon/latest/rayon/

まずは素直に実装してみる

素直に、dp[i][j]を $i$ 人の子供に $j$ 個のキャンディを配る時の組み合わせだと考えると、 $i$ 人目の子供が受け取れるキャンディの数は、0から $a_i$ 個の幅なので、dp[i-1][j]からdp[i-1][j+a[i]]の値を加えたものになる。実装では結局前の結果しか不要なのでテーブルではなく、2つのベクターで実装しており、さらに「配る動的計画法」とも呼ばれているdp[i-1]の方をループで回す方法で実装している。

candies_straightforward.rs

use proconio::input;
use std::cmp::min;

const MOD: u64 = 1000000007;

fn main(){
    input! {
        children: usize, candies:usize,
        mut a: [usize; children]
    }
    let candies = candies + 1;
    let mut dp = vec![1u64; 1];
    for i in 0..children {
        let mut next_dp = vec![0; min(candies, dp.len() + a[i])];
        for j in 0..dp.len() {
            for k in j..min(j + a[i] + 1, candies) {
                next_dp[k] = (next_dp[k] + dp[j]) % MOD;
            }
        }
        dp = next_dp;
    }

    if dp.len() < candies {
        println!("0");   
    } else {
        println!("{}", dp[candies - 1]);        
    }
}

実行速度の測定

リストを見ると3重ループとなっており、一番外は子供の数 $N$ 、内側の2つはキャンディの数 $K$ に比例している^[1]ので、計算量は $O(NK^2)$ となる。問題の条件としては $1\leq N \leq100$ , $0\leq K \leq 10^5$ なので、 $K^2$ の部分がやばい感じ。

実際に $N=100$ として $K$ を変化させて計測すると以下のような感じになる。25,000あたりから1秒を越えて35,000あたりで実行時間制限の2秒を越えている。入力部分を入れずに時間制限を越えるので、この実装はNGであろう。ベクタaをソートしたり小技はありそうな気はするが、実行時間に大差はないと思われる。

累積和を使う

内側のkのループを外せば計算量が $O(NK)$ になるので、外す方法を考える。ここはjからa[i]の間にdp[j]を加えていくという処理になっている。これをjの位置に+dp[j]をa[i]+1の位置に-dp[i]を設置して、最後に累積和を取るという形に変更してみる。つまり $K^2$ を $2K$ に変更する。

candies_straightforward.rs

for k in j..min(j + a[i] + 1, candies) {
  next_dp[k] = (next_dp[k] + dp[j]) % MOD;
}

ループでの実装

実際に実装すると以下のようになる。kのループがなくなる代わりにjのループが2つ並んでいる。

candies_loop.rs

use proconio::input;

const MOD: i64 = 1000000007;

fn main(){
    input! {
        children: usize, candies:usize,
        mut a: [usize; children]
    }
    let candies = candies + 1;
    let mut dp = vec![0i64; candies];
    dp[0] = 1;
    for i in 0..children {
        let mut next_dp = vec![0; candies];
        for j in 0..candies {
            next_dp[j] = dp[j];
            if j + a[i] + 1 < candies { 
                next_dp[j + a[i] + 1] = (MOD + next_dp[j + a[i] + 1] - dp[j]) % MOD;
            }
        }
        for j in 1..candies {
            next_dp[j] = (next_dp[j] + next_dp[j - 1]) % MOD;
        }
        dp = next_dp;
    }
    println!("{}", dp[candies - 1]);
}

Rayonを使った並列化

そもそも累積和を取る前のnext_dpは、「dp」から「a[i]+1シフトしたdp」を引くだけで求められる。rotateはあるがshiftはなかったのでextend_from_sliceを使ってシフトしたdpを作成した。さらに、iter()をpar_iter()に変更するだけで並列化してくれる素晴らしいcrateであるrayonを使って並列化までやってみた。

candies_rayon.rs

use proconio::input;
use rayon::prelude::*;

const MOD: i64 = 1000000007;

fn main(){
    input! {
        children: usize, candies:usize,
        mut a: [usize; children]
    }
    let candies = candies + 1;
    let mut dp = vec![0i64; candies];
    dp[0] = 1;
    for i in 0..children {
        let mut shift_dp = vec!(0; a[i] + 1);
        shift_dp.extend_from_slice(&dp[0..(candies - (a[i] + 1))]);
        let mut next_dp: Vec<i64> = dp.par_iter()
            .zip(shift_dp.par_iter())
            .map(|(b, e)| (MOD + b - e) % MOD).collect();
        for j in 1..candies {
            next_dp[j] = (next_dp[j] + next_dp[j - 1]) % MOD;
        }
        dp = next_dp;
    }
    println!("{}", dp[candies - 1]);
}

処理速度の計測

計測結果は以下の通り。iterはrayonのコードでpar_iter()ではなく、iter()を使ったもの、つまり、並列化をしていないものである。課題の $0\leq K \leq 10^5$ の範囲では問題なく、 $10^6$ まで延ばしてみたが、差はさほど大きく無い。rayonが一番早いがiterと比較すると並列化の恩恵は小さく、shift_dpの準備や累積和の処理が大きい印象。並列化の効果というのは中々でにくいなと感じる。

ただ、どの方法も100ms以下で終了しており、課題の解答としては問題ない。

さらに高速化

並列化の効果が低いとすると、最初から累積和をとればいいのではないか？ということで書き直してみた。

candies_loop_2.rs

dp[0] = 1;
for i in 0..children {
    let mut next_dp = Vec::with_capacity(candies);
    next_dp.push(dp[0]);
    for j in 1..candies {
        if j < a[i] + 1 {
            next_dp.push((dp[j] + next_dp[j - 1]) % MOD);
        } else {
            next_dp.push((dp[j] + next_dp[j - 1] + MOD - dp[j - (a[i] + 1)]) % MOD);    
        }
    }
    dp = next_dp;
}

結果としては、Rayonで並列化したものよりも速い。コードもすごくシンプルで速い。何故最初からこれがでてこなかったんだろうかと思う。

脚注

$0\leq a_i \leq K$ なので $a_i$ も $K$ に比例する ↩︎