Web系エンジニアのアウトプット練習場

エンジニアリングのことはもちろん、全然関係無い話もします。

【AtCoder Beginner Contest 132】感想戦

A, B, Cの3完。パフォは625で過去ワーストを100以上更新して凹んでいます。

f:id:h-sakano:20190629233130p:plain

今回詰まったD問題のハイライト。

「ふむ、N-K+1個の隙間からi箇所選んで、更にK個のボールをiこのグループに分ける組み合わせを掛けたらそれが答えだな」 「よし!フェルマーの小定理使うやつや!」

一瞬でおおまかな方針がたち、ぬか喜び。

K個のボールをiこのグループに分ける組み合わせってどうやって求めるっけ?」

組み合わせの知識を呼び戻すために「グループ分け 組み合わせ」でぐぐってみると、「9人の生徒を3人ずつのグループに分ける分け方は何通り?」とかグループの人数が決まっているものしかなくて焦る。

「もうインターネッツの力には頼らん。DPで解いたる。」と以下のコードが爆誕。

N, K = map(int, input().split())
MOD = 10 ** 9 + 7


def comb(n, k, p):
    """power_funcを用いて(nCk) mod p を求める"""
    from math import factorial
    if n < 0 or k < 0 or n < k:
        return 0
    if n == 0 or k == 0:
        return 1
    a = factorial(n) % p
    b = factorial(k) % p
    c = factorial(n-k) % p
    return (a*power_func(b, p-2, p)*power_func(c, p-2, p)) % p


def power_func(a, b, p):
    """a^b mod p を求める"""
    if b == 0:
        return 1
    if b % 2 == 0:
        d = power_func(a, b//2, p)
        return d*d % p
    if b % 2 == 1:
        return (a*power_func(a, b-1, p)) % p


memo = [[-1 for _ in range(K + 1)] for _ in range(K + 1)]


def dp(n, k):
    # n個のボールをk個のグループに分ける(空のグループを許容)
    if n <= 0 or k <= 1:
        return 1
    if memo[n][k] >= 0:
        return memo[n][k]
    memo[n][k] = 0
    for i in range(n, -1, -1):
        memo[n][k] += dp(n - i, k - 1) % MOD
    return memo[n][k]


for i in range(1, K + 1):
    print(comb(N - K + 1, i, MOD) * dp(K - i, i) % MOD)

O(n^3)のコードができて、めでたくTLEを喰らいましたとさ。

〜完〜

どうすればよかったのか

K個のボールをiこのグループに分ける組み合わせ =  _{K-1}C_{i-1}

公式の解説がわかりやすかったです。

K個のボールを並べて、K − 1個のボールとボールの間からi − 1箇所選び、そこで切り分けるとi箇所に分かれる。各部分は最小1個のボールを含んでおり、ボールの総数は K個になる。

なるほど!参りました!

まぁ今回はフェルマーの小定理が使えただけでもよしとしましょう(負け惜しみ)

今回の凡ミス

A問題提出時に言語選択が数時間前にちょっと練習してたRustのままになっててCE食らった。