🪨

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

2023/04/19に公開

Rust

AtCoder

動的計画法

tech

はじめに

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

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

AからZまで問題が設定されているが、今回はKのStones。石を交互に取っていって、取れなくなった方が負けというゲームの解放である。

利用したライブラリ

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

完成したコード

実装に難しい所はない。

stones.rs

use proconio::input;

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

    let k = k + 1;
    let mut dp = vec![false; k];
    for i in 0..k {
        if dp[i] == false {
            for j in 0..n {
                if i + capture[j] < k { dp[i + capture[j]] = true; }
            }
        }
    }
    if dp[k - 1] == true {
        println!("First");
    } else {
        println!("Second");    
    }
}

アルゴリズム

石の残りが0であれば「負け」。「負け」の状態になるように石が取れるなら「勝ち」。「勝ち」の状態になるようにしか石が取れないなら「負け」。ということで、石の数が0の状態から、「負け(false)」の状態に移動できる石の数を「勝ち(true)」として、石の数が $K$ 個になった状態で「勝ち」であれば、今から石をとるのは先行なので、先行の勝ち。異なれば後攻の勝ち。

先行と後攻を意識して考えたくなるが、意識しなくて解けるところが面白いところだろうか？状態（石の数）で勝ち負けが決まってしまうというシンプルなゲームだからであるが。

ビット演算で解いてみる

出題されている問題において、石の数は最大 $10^5$ 個なので適合しないが、ビット演算を使って実装してみる。今回はbool型の配列しか使っていないので、それを0, 1のビットで表現する。ビットが0(負け)である時に、その状態に遷移可能な位置を示すcapture_maskをシフトさせて論理和orを取れば同じことが実現できる。

stones_bit.rs

fn main(){
    input! {
        n: usize, k:usize,
        capture: [i32; n]
    }

    let mut dp: u128 = 0;
    let mut capture_mask: u128 = 0;
    for i in 0..n {
        capture_mask = capture_mask | (1u128 << (capture[i] - 1));
    }
    
    dp = dp | capture_mask;
    for i in 0..k {
        if dp & (1u128 << i) == 0u128 {
            dp = dp | (capture_mask << (i + 1)) 
        }
    }

    if dp & (1u128 << (k - 1)) == 0u128 {
        println!("Second");        
    } else {
        println!("First");
    }
}

Rustにu128型があったので、128ビット列を使っている。0ビット目を石の数が1個の場合としているので、 $K$ が128までの問題であれば解ける。汎用性や拡張性に乏しいので難しいが、こんな実装方法もあるというサンプルである。

計算量の測定

$K\leq 128$ の問題までしか解けないので比較が難しいが、1000回同じ問題を繰り返し解いて、その計算時間を1000分の1にすることで測定した結果を示す。当然であるがビット演算の方が $j$ のループが無くなっているので早い。とれる石の数の集合 $A = \{2, 3, 5, 7\}$ で固定してある。ビット演算の場合は、 $A$ のサイズ $N$ に依存しないが、比較できる程大きな問題が解けない問題がある。