なんとか仕事は山を越えました。明日からはきちんと定時に帰って、記事を書く頻度を上げられるようにしたいです（フラグ）。

最近はRustもだいぶスラスラ書けるようになった気がします（ただし、競技プログラミングのみ）。制限時間がシビアな高得点問題をRustで解く選択肢もそろそろ見えてきました。

ABC132-F

入力: 正の整数N, K
出力: （隣り合う項の積がN以下であるような長さKの正の整数列）の個数

https://atcoder.jp/contests/abc132/tasks/abc132_f

解き方

ポイントは正の整数X, A, B について、X = AB ならば min(A, B) <= sqrt(X) が成り立つことです。

P = floor(sqrt(N)), Q = floor(N / M) とします。 [1, N] の整数から成る集合を以下のP + Q - 1 個の集合に分割します。

G_1: 1のみ含む集合
...
G_{P}: P のみ含む集合
G_{P + 1}: floor(N / i) = Q - 1 を満たす整数iの集合
G_{P + 2}: floor(N / i) = Q - 2 を満たす整数iの集合
...
G_{P + Q - 1}: floor(N / i) = 1 を満たす整数iの集合

このとき、以下が成り立ちます。

i <= P ならば |G_i| = 1
i > P ならば |G_i| = floor(N / (P + Q - i)) - floor(N / (P + Q - i + 1))
G_i に属す要素g_iと G_j に属す要素g_j について、i + j <= P + Q ならば g_i g_j <= 2N

長さがL、かつ、隣接して並んでいるどの2つの整数の積もN以下、かつ、列の最後の要素が G_i に属すような列の個数を dp[L][i] とします。以下が成り立ちます。

L = 1 ならば dp[L][i] = |G_i|
L >= 2 ならば dp[L][i] = |G_i|(dp[L - 1][1] + ... + dp[L][P + Q - i])

求める解は dp[L][1] + ... + dp[L][P + Q - 1] となります。

RustでDPするには

値の変更とか面倒そうですが、意外と簡単にDPできます。例えば、要素数Nの1次元配列を用いてDPしたい場合、

let mut dp = vec!(0; N);
dp[0] = 1;
dp[1] = 1;
for i in 2..N {
  dp[i] = dp[i - 1] + dp[i - 2];
}

のように書くだけです。 vec!(x; N) で各要素がxである長さNのVecが得られます。 Vecの値は書き換えたいので、 mut を忘れないように気をつけましょう（忘れるとコンパイラに怒られます）。

注意: Vecの添字にはusize型を指定する

dp[i] のようにVecのi番目の要素を取得する際、iはusize型である必要があることに注意です。 Nをusizeの値とすると、 (0..N) が生成する整数はusizeとなるので、整数が非負ならばusize型として宣言したほうがいいかもしれません。

コード

fn main() {
    let (N, K) = get!(usize, usize);
    let MOD: usize = 1000000007;

    // floor(N / i) の値に応じてグループ分けし、
    // グループ単位で条件を満たす列の個数を動的計画法で求める
    let P = (N as f64).sqrt() as usize;
    let Q = N / P;
    let R = P + Q;
    let C = (0..R).map(|i|
        if i <= P {
            1
        } else {
            N / (R - i) - N / (R - i + 1)
        }
    ).collect::<Vec<_>>();
    let mut dp = vec![vec![0; R]; K];
    for i in 0..R {
        dp[0][i] = C[i];
    }
    for j in 1..K {
        let mut acc = 0;
        for i in (1..R).rev() {
            acc = (acc + dp[j - 1][R - i]) % MOD;
            dp[j][i] = (C[i] * acc) % MOD;
        }
    }
    let ans: usize = dp[K - 1].iter().sum::<usize>() % MOD;

    println!("{}", ans);
}

ちなみに上記と同等のPythonコードをPyPyとして提出したらTLEでした。一方で、上記Rustコードは45ms.　圧倒的ではないか！ Pythonは縛りプレイと呼ばれている理由が理解できますね。

024号文書

主にプログラミング

PythonistaがRustはじめました#008 -- 可変Vecで動的計画法

ABC132-F

解き方

RustでDPするには

注意: Vecの添字にはusize型を指定する

コード