【Python3】 AtCoder Beginner Contest 156 D – Bouquet

二項係数 $_n \mathrm{C} _r$ の高速計算の実装に手こずった。

サトゥー

$n$ が莫大で $r$ が比較的小さいときに有効な二項係数の計算です。

あかりさんは $n$ 種類の花を $1$ 本ずつ持っています。
あかりさんは、これらの花から $1$ 本以上を選び、花束を作ろうとしています。
ただし、あかりさんは $a$ と $b$ の $2$ つの数を苦手としていて、いずれかと一致するような本数の花からなる花束は作ることができません。
あかりさんが作ることのできる花束は何種類あるでしょうか。$(10^9+7)$ で割った余りを求めてください。
ここで $2$ つの花束は、一方では使われているが、 もう一方では使われていない種類の花があるとき、別の種類の花束であるとみなします。

制約
・ $2 \leq n \leq 10^9$
・ $1 \leq a < b \leq min(n, 2 \times 10^5)$

https://atcoder.jp/contests/abc156/tasks/abc156_d

考察は単純で、 $n$ 種類の花から作れる花束のすべては $2^n-1$ 通りで、ここから $a$ 本、 $b$ 本で花束を作る場合を除けばよく、

$$ (2^n \, – \, 1) \, – \, _n \mathrm{C} _a \, – \, _n \mathrm{C} _b$$

となることがわかります。

ここから、 $ \mathrm{mod} \, 10^9 + 7 $ での組み合わせ数の計算をするわけですが、下の記事で見ているような逆元テーブルを用意してちまちま計算していく方法だと今回は $n$ がデカすぎて詰みます。

そこで今回は、 $_n \mathrm{C} _r$ における $r$ が比較的小さい値になっていることに着目し、

$$ _n \mathrm{C} _r = \frac{n \cdot (n-1) \cdot …… \cdot (n-r+1)}{1 \cdot 2 \cdot …… \cdot r} $$

として考えてみます。

分子 ($numer = n \cdot (n-1) \cdot …… \cdot (n-r+1)$ ) は普通に mod を取りながら計算するとして、 分母 ( $denom = r!$ ) の逆元は次のように考えてみます。

フェルマーの小定理を用いた逆元計算

$p = 10^9 + 7$ は素数なので、フェルマーの小定理より以下の式が成り立ち、結局整数 $m$ の $\mathrm{mod} \, p$ における逆元は $m^{p-2}$ と表現できる。

$$ m \times m^{p-2} \equiv 1 \, ( \mathrm{mod} \, p )$$

Python では関数 pow() を使うことで、この逆元を計算できます。具体的に整数 $m$ に対して逆元は次のように表現できます。

pow(m, p-2, p)

こんな感じで $_n \mathrm{C} _a$ と $_n \mathrm{C} _b$ を計算すればなんとかなりそう。基本的な $_n \mathrm{C} _r$ の実装は 【Python】組み合わせ(nCr) 計算の高速化 – Qiita を参考にしつつ、 mod の要素を取り入れてみます。

と、ここまでをふまえて、回答例は以下のようなコードになります。

from functools import reduce

def cmb(n, r, p):
    r = min(n - r, r)
    if r == 0:
        return 1
    numer = reduce(lambda x, y: (x*y)%p, range(n, n - r, -1))
    denom = reduce(lambda x, y: (x*y)%p, range(1, r + 1))
    return (numer * pow(denom, p-2, p)) % p

p = 10 ** 9 + 7
n, a, b = map(int, input().split())

ans = pow(2, n, p) - 1
ans -= cmb(n, a, p) + cmb(n, b, p)
print(ans % p)


この記事を書いた人

サトゥー

東大学際情報学府M1。情報科学と教養の海に溺れています。面白いことをやるのがすきです。