🖥️

Rustで動的計画法：Digit Sum

2023/05/20に公開

Rust

AtCoder

動的計画法

tech

はじめに

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

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

AからZまで問題が設定されているが、今回はSのDigitを解いてみる。 $K$ 以下で各桁の数値の合計が $D$ の倍数である整数の数を数え上げる問題。解き方は難しいが、コーディングはそうでもないので、テストコードの書き方でも学んでみる。

利用したライブラリ

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

プログラム自体のテストを行うためにcli_test_dirを利用。

完成したコード

$K$ の上位 $i$ 桁を $K_i$ とした時、 $K_i$ より小さくて、 $d$ で割った余りが $r$ となる整数の数をdp[i][r]として、動的計画法を動かせばいい。 $i$ から $i + 1$ の導出については、 $i$ 桁目の数値を $a_i$ とすると、

\sum_j^{i+1} a_j \bmod d = \left(\left(\sum_j^i a_j \bmod d\right) + a_{i+1}\right) \bmod d

となることを利用して求める。さらに上位 $i$ 桁目まで $K_i$ に等しいならば、 $i+1$ 桁目は $a_{i+1}$ 以下である必要があるので、上位 $i$ 桁目まで $K_i$ の場合と、 $K_i$ 未満である場合とで異なるループを回す必要がある。

digit_sum.rs

use proconio::input;
use proconio::marker::Chars;

const MOD:u32 = 10u32.pow(9) + 7;

macro_rules! I {
    ( $i:expr, $j:expr, $d:expr ) => ( (($i + $j) % $d) as usize )
}

fn main(){
    input! {
        k: Chars,
        d: u32,
    }

    let mut dp = vec![0u32; d as usize];
    
    let mut digit = k[0].to_digit(10u32).unwrap();
    for i in 0..digit {
        dp[(i % d) as usize] += 1;
    }
    let mut digit_sum_k = digit % d;

    for n in 1..k.len() {
        let mut next_dp = vec![0u32; d as usize];
        digit = k[n].to_digit(10u32).unwrap();

        for i in 0..digit {
            next_dp[I!(i, digit_sum_k, d)] = (next_dp[I!(i, digit_sum_k, d)] + 1) % MOD;
        }

        for i in 0..10 {
            for j in 0..d {
                next_dp[I!(i, j, d)] = (next_dp[I!(i, j, d)] + dp[j as usize]) % MOD;
            }
        }
        dp = next_dp;
        digit_sum_k = (digit + digit_sum_k) % d;
    }
    let ans = if digit_sum_k == 0 { dp[0] } else { dp[0] - 1 };
    println!("{}", ans)
}

charの配列

入力 $K$ の上限が、 $10^{10000}$ に設定してあるので、uint64でもオーバーフローしてしまう。このため、文字列として取得する。proconioだと、marker::Charsを使うと、charのvectorとしてデータが読み取られる。

input! {
    k: Chars,
    d: u32,
}

使う時はto_digit()で数値化する。結果はOption型なので、unwrap()する。

let mut digit = k[0].to_digit(10u32).unwrap();

macro

dpの更新の部分をちゃんと書くと以下のようになる。

next_dp[((i + digit_sum_k) % d) as usize] = (next_dp[((i + digit_sum_k) % d) as usize] + 1) % MOD

なんとなく添え字がごちゃごちゃするので、マクロI!を定義してみた。Rustのマクロは引数のパターンマッチングを並べる形で記述する。今回は、3つの数式を引数としたマッチングを定義した。
見やすくなったかは微妙かもしれないが。

macro_rules! I {
    ( $i:expr, $j:expr, $d:expr ) => ( (($i + $j) % $d) as usize )
}

テスト

AtCoderの問題には、入力例と出力例（つまり解答）が載っているので、それでプログラムがちゃんと実装されているかを確認できる。Rustには今時の言語っぽくテスト用の機能もついているので、それを使っていく。

まず、testsディレクトリ以下のファイルはテスト用だと思われるので、その下にファイルを置く。
そこにファイルを置いてテスト用の関数を書く。#[test]がついている関数がテスト時に順番に呼び出される。

今回は、コマンドラインのプログラムに対してstdinから何らかのデータを入力して、その結果をstdoutから取得するので、そのためのライブラリcli_test_dirを使う。

tests/samples.rs

use cli_test_dir::*;
const BIN: &str = "./digit_num";

#[test]
fn sample3() {
    let testdir = TestDir::new(BIN, "");
    let output = testdir
        .cmd()
        .output_with_stdin(r#"98765432109876543210
        58"#)
        .tee_output()
        .expect_success();
    assert_eq!(output.stdout_str(), "635270834\n");
    assert!(output.stderr_str().is_empty());
}

内容としてコードを見ればそのままという感じ。BINはCargo.tomlで指定したpackageのnameを書けばいい。

r#" "#はraw stringで、改行も含めてそのままの文字列を示す。#の数は任意なので文中に"#が入るなら、#を増やせばいい。（ないなら無くてもいいのでこの場合はr" "rでも良い）

cargo textで結果が得られる。

% cargo test
running 1 test
635270834
test sample3 ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.06s

過去の記事

Rustで動的計画法の実装: