昨日(2019/11/16)のAtCoder Beginner ContestのD -Kightの解説です。
主に、$nCk$ $\pmod{10^9+7}$の計算方法について書きます。
考え方
問題の$X$, $Y$の制約が、
・$1 \leqq X \leqq 10^6$
・$1 \leqq Y \leqq 10^6$
であることから、DPやBFS, DFSではTLEになってしまうことを察する必要があります。
（どれもO(NM)のオーダーになってしまいます。）
※ちなみにAtCoderの計算量の限界は、$1.0\times10^7$ 程度と言われています。
上のアルゴリズムでダメなら、何か一工夫が必要ということなるので、ここで解説に出てくるような考え方が思いつけば、あとは $nCk$ $\pmod{10^9+7}$を求めるだけの問題となります。
nCkの求め方・考え方
今回求める必要があるのは、$nCk = \frac{n!}{k!(n-k)!}$を、$10^9+7$で割ったあまりです。
※今回の対象は、$1 \leqq k \leqq n \leqq 10^7$程度の値の場合です。
毎度ながらけんちょんさんの記事を参考に、自分の思い出し用にまとめます。
求める上でポイント（ややこしい点）が４つあるので、順番に見ていきましょう。
Ⅰ. nCkの計算の際、掛け算のたびに余りを取る
これはmod計算の基本ですね。オーバーフローを起こしてしまうので、積のたびに毎回余りを取ります。今回の式の場合、
$\frac{n!}{k!(n-k)!} = n \times (n-1) \times ... \times 1 \times (\frac{1}{k} \times \frac{1}{k-1} \times ... \frac{1}{1}) \times (\frac{1}{n-k} \times \frac{1}{n-k-1} \times ... \frac{1}{1})$
となるのですが、各要素 $n-1, \frac{1}{n-k-1}$ を計算する際に、$mod$ $10^9+7$を取るということです。
ただしここで次のⅡの疑問が生じます。
Ⅱ. mod p における割り算とは何か？
掛け算の時に毎回$mod$ $10^9+7$を取るのはわかったが、$\frac{1}{n-k-1}$は実質割り算であり、どう計算すれば良いのか、という話です。
これは、けんちょんさんの記事、
がわかりやすいです。
**「3. 割り算 a ÷ b」**のところに、mod p の世界の割り算とは何か、が詳しく書いてあります。
結論だけ書くと、
$\mod{p}$において、$b$をかけると$1$になる数を、$b^{-1}$と表記すると
$ a \div b \equiv a \times b^{-1} \pmod{p}$
となります。つまり、$b^{-1}$を求める（そこにaを掛ける）ことが、modの世界の割り算だということです。
$b^{-1}$のことを、mod $p$における$b$の逆元といいます。
Ⅲ. 逆元をどうやって求めるか？
では、Ⅱで出てきた逆元はどのように求めればよいのでしょうか。
これも、けんちょんさんの記事、
に詳しく説明が書いてあります。
個人的には、「1-5. 逆元漸化式のもう 1 つの導出方法」の説明が特にわかりやすかったです。
こちらも結論だけ書くと、
$b^{-1} \pmod{p}$
を計算するとはすなわち、
$ b^{-1} \equiv -(p$%$b)^{-1} \times (p/b) \pmod{p} $
の右辺を計算することだと言えます。
※ここで、$p$%$b$は$b$より小さくなることに注意
実装上は、$b$未満の逆元の要素が入ったリストをあらかじめ用意しておくことで、以下のように書くことできます(python)。
# invは1～(b-1)^-1までの逆元の値が入ったリスト
inv[b] = - inv[MOD % b] * (MOD // i) % MOD
ちなみにけんちょんさんの実装例では、以下のようになっていますが、右辺にMODを足して正の数にしているだけなので、同じことです。
inv[b] = MOD - inv[MOD % b] * (MOD // i) % MOD
例）
・ $3 \pmod{11} \equiv 3 + 11 = 14 $
・ $-3 \pmod{11} \equiv -3 + 11 = 8 $
Ⅳ. nCkを効率良く求めるには？
上のⅠ～Ⅲがわかれば、あと少しです。
求めたいのは$p=10^9+7$として、以下の式の、$\pmod{p}$です。
$\frac{n!}{k!(n-k)!} = n \times (n-1) \times ... \times 1 \times (\frac{1}{k} \times \frac{1}{k-1} \times ... \frac{1}{1}) \times (\frac{1}{n-k} \times \frac{1}{n-k-1} \times ... \frac{1}{1})$
これを求めるためには、
(1)$1～n$までの階乗のMODのリスト
[$1!\pmod{p}$, $2!\pmod{p}$, ..., $n!\pmod{p}$]
(2)逆元の階乗リスト
[$\frac{1}{1!}\pmod{p}$, $\frac{1}{2!}\pmod{p}$, ..., $\frac{1}{n!}\pmod{p}$]
さえあれば、簡単に求まることがわかります。
(1)は
$n!\pmod{p} \equiv (n-1)!\pmod{p} \times n$
というように、前の要素を用いて簡単に求めることができます。
(2)に関しても、
$\frac{1}{n!}\pmod{p} \equiv \frac{1}{(n-1)!}\pmod{p} \times \frac{1}{n}(=n^{-1})$
という式で求められますが、Ⅲで学んだように、$n^{-1}$を求めるには、$0$～$n-1$までの逆元のリストが必要です。
したがって、(1)、(2)に加えて、
(3)逆元のリスト
[$\frac{1}{1}\pmod{p}$, $\frac{1}{2}\pmod{p}$, ..., $\frac{1}{n}\pmod{p}$]
の３つがそろえば、$O(1)$のオーダーで$nCk$の計算をすることができます。
以下は今回の問題の解答例です。
def cmb(n, k, mod, fac, ifac):
"""
nCkを計算する
"""
k = min(k, n-k)
return fac[n] * ifac[k] * ifac[n-k] % mod
def make_tables(mod, n):
"""
階乗テーブル、逆元の階乗テーブルを作成する
"""
fac = [1, 1] # 階乗テーブル・・・(1)
ifac = [1, 1] #逆元の階乗テーブル・・・(2)
inverse = [0, 1] #逆元テーブル・・・(3)
for i in range(2, n+1):
fac.append((fac[-1] * i) % mod)
inverse.append((-inverse[mod % i] * (mod//i)) % mod)
ifac.append((ifac[-1] * inverse[-1]) % mod)
return fac, ifac
X, Y = map(int, input().split())
# 後の、n, mの大小を一定にするため
if X > Y:
X, Y = Y, X
dist = X + Y
if dist % 3 != 0:
print(0)
else:
# (+1,+2) の移動の回数を n、(+2,+1) の移動の回数を m とする)
# total = n + m
# 以下の式は、2n + m = X, n + 2m = Y の連立方程式を解いた結果
total = int((X+Y) / 3)
n = X - total
# n < 0 もしくは m < 0 の場合、たどり着けない
# Y > X なので、m > n
# n < 0 → 2X - Y < 0
if Y > 2*X:
print(0)
else:
MOD = 10**9 + 7
fac, ifac = make_tables(MOD, total)
ans = cmb(total, n, MOD, fac, ifac)
print(ans)
備考
今回は、割る数が$10^9+7$だから問題なく計算できましたが、どんな数でも良いわけではないようです。
割る数$p$と$nCr$の関係は、
- $p$は素数
- （今回の実装に関しては）$p > n$
という条件が必要になります。