組み合わせnCrの総和 計算の高速化及び逆元について(Python編)

はじめに

よろしくお願いします。
競技プログラミングをご存じでしょうか。
毎週土日や平日の夜中に様々なサイトで行われています。(atcoder ,codeforces, topcoder等々　拙者はatcoderのみですが)

<wikipediaより

競技プログラミングでは、参加者全員に同一の課題が出題され、より早く与えられた要求を満足するプログラムを正確に記述することを競う。
コンピュータサイエンスや数学の知識を必要とする問題が多く、新卒学生の採用活動などに使われることもある

今回は競技プログラミングによく使う、組み合わせについて書いていきたいと思います。

初心者向けであるのでつよつよ競プロerさんには絶対に必要ありません。
読まないでください（笑）（間違いがあればご指摘ください）

ちなみにPython編とありますがPythonしか書きません。

今回の議題

競技プログラミングにおいて組み合わせの問題であったり、ここに帰着させる問題が時折登場します。

例えば

nまでの数字(被りなし）から1個、2個、...n個それぞれ取り出すときの組み合わせの総和を 998244353 で割った余りを求めなさい。
n<=10**5
最初にコードを見たい方

例　n=10
1個の時　10 C1=10
2個の時　10 C2=45
3個の時　10 C3=120
...
10個の時　10 C10=1
解　1023

＊難しいものに取り組みたい方用の問題

組み合わせの公式については以下の通りである。

nCr = \frac{N!}{(N-R)! R!}

以下に初心者であったら自分が書くであろうコードを示す。

wrongcode.py

import math
mod=998244353
n=10000

#関数部分
def comb(N,R):
    return (math.factorial(N)//math.factorial(N-R)//math.factorial(R))%mod
#
p=0
for i in range(1,n+1):
    p+=comb(n,i)
    p%=mod
print(p)

math.factorial は階乗を求める関数である。

上のコード実際に回してみよう。
n=10では余裕をもって通る。
n=1000も無事ACである。
しかしn=10000つまり10**4 の時に制限時間以内に終わらずTLEとなってしまう。

少し上級者になれば

実はこのコードには一番簡単な解決法が一つある。
それはメモ化である。

上記のコードにおいて
組み合わせの式には階乗が何度も用いられている。
よって、前計算してどこかに記憶して取り出せればよいのではないかと思い付く。

解決案

#関数部分
n1=[1]+list(range(1,n+1))
for i in range(1,n+1):
    n1[i]*=n1[i-1]
def comb(N,R):
    return (n1[N]//n1[R]//n1[N-R])%mod
#

これで無事解決！！！

.....と行きたいところであるがそうもいかない。
上記同様、n=10000つまり10**4 の時に制限時間以内に終わらずTLEとなってしまう。

なぜだろうか？
メモ化が機能していない？
どこかでコードを間違えている？
自分も初めは悩みこんでしまった。

とても簡単な問題であるにもかかわらず
実は初心者にとってかなり厄介な問題を抱えているのである。
なので、今回はこの問題を取り扱っていきたい。
初心者の助けになればうれしいと思う。

どこが厄介なのか

なぜだろうか？
二つの点について順序だてて話したい。

・あまりに関する問題

まず、998244353で割ったあまりに問題点が存在している。

余りについてはよく使われる基本なので知っておいて損はないが
知っていなくても問題ないので隠しタブとした

なぜ余りを使っているのか

競プロではなぜ余り(mod = 1000000007 ,998244353)を取るのか？

この数字は、十分に大きい、素数であるというところがミソなのである。

PCでは32bitや64bit数字で計算することが多い。
それ以上になる場合は多倍長整数が使われる。
これは言語によってなかったり、計算に時間がかかるという特徴がある。

そこで余りを取れば計算量を落とし、解を得ることがが可能となる。
この数字であるのはmod以下の数字では、足したときに32bitを超えない。
そして掛け合わせた時に64bitを超えないという特徴があるからだ。

また、素数とする理由としては常にその数字以下の数字において逆元が存在するという理由がある。
逆元について今から書いていこうと思う。

参考文献　Mond 競技プログラミングの問題によく出てくる1,000,000,007とか998244353は、どこから出てきた数字なのでしょうか？

この問題において998244353を超えないで計算を行いたい。
しかしコードを見ていただくと

先ほどのコード一部

import math
math.factorial(N)

この階乗の部分に問題がありそうである。
なぜなら、

N=10　3628800
N=100 : 93326215443944152681699238856266700490715968264381621468592963895217599993229915608941463976156518286253697920827223758251185210916864000000000000000000000000

とてつもなく大きい数字なっている。
n=100の時よくバグってなかったとほめたいくらいである。
ここに改善点の余地がある。
そこで掛けるたびに余りを取るという解決法が思いつく。

wrong2.py

#関数部分
for i in range(1,n+1):
    n1[i]*=n1[i-1]
    n1[i]%=mod
def comb(N,R):
    return n1[N]//n1[R]//n1[N-R]
#

すると次にはもう一つ問題が発生する。

・割り算における問題

まず、上記の方法を取り入れた時にバグる例から考えてみよう。
これは計算を行うとわかるがn=13以上になるとn!>998244353の数字となる。
よって今回は13C2を求めることでバグについて考えてみよう。

以上の計算を行うと
関数内で

13!は6227020800 より余りを取ると　分子 = 237554682
分母 = 11! * 2! = 79833600
13C2=0

よって間違った答えが出てきてしまう。
余りを取りたくないなと考えるに至ってしまう。
ほんとにどうしようか？

解決するには

逆元という手法である。

逆元とは

a≡\frac{1}{p} \;  (mod\:M)

上記のような式を用いることでpで割るのではなくaを掛けることで求めることが可能になる。

この手法について述べる。
そしてその実装方法二つについて詳しく述べていこうと思う。

• フェルマーの小定理
• 拡張ユークリッドの互除法

を用いた実装について述べる

逆元について

本来であれば詳しく書きたいところであるが
実はあまり詳しくないので「逆元とは何か」と「逆元の求め方」について公式を用いて書いていきたいと思う。

逆元とは何か

逆元とは何か?
例を用いて話していく。
まずは底を7とした時で考えていきたい。

0 ≡ 0\;  (mod\:7)

1 ≡ 1\;  (mod\:7)

2 ≡ 2\;  (mod\:7)

3 ≡ 3\;  (mod\:7)

4 ≡ 4\;  (mod\:7)

5 ≡ 5\;  (mod\:7)

6 ≡ 6\;  (mod\:7)

7 ≡ 0\;  (mod\:7)

8 ≡ 1\;  (mod\:7)\\

ではここで8÷2がしたくなったと想定した時、本来であれば4となります
しかし余りを取っていたときは8 ≡ 1となり、答えが0.5と間違った数字となります。

そこで2の逆元、つまり1/2と等しいものを算出します。
後で説明する計算を行うとmod7における　1/2の逆元は4となります。

8 ≡ 1\;  (mod\:7)

\frac{1}{2} ≡ 4\;  (mod\:7)

1÷2 ≡ 1*\frac{1}{2} ≡ 1*4 ≡ 4 \;(mod\:7)

このようにすると求められます。

別の数字で試してみましょう。
20÷5がしたくなったと想定した時、本来であれば4となります(mod 17)
1/5のmod 17における逆元は7であるため

20 ≡ 3\;  (mod\:17)

\frac{1}{5} ≡ 7\;  (mod\:17)

3÷5 ≡ 3*\frac{1}{5} ≡ 3*7 ≡ 21 ≡ 4 \;(mod\:17)\\

(2個でめんどくさくなったのでいいですかね・・・)
このように、逆元を求めると計算がスムーズに行うことが可能になります。
次に実装方法を述べる。

逆元の求め方

これは自分の知っている範囲で2つある。

• フェルマーの小定理
• 拡張ユークリッドの互除法

両方とも理論の説明と実装を行う

フェルマーの小定理の理論及び実装

フェルマーの小定理
M が素数で，a が M の倍数でない正の整数のとき

a^{M-1} ≡ 1\;  (mod\:M)\\

証明については省略する。

この式を変形すると

a^{M-1} ≡ 1\;  (mod\:M)\\
a^{M-2} ≡ \frac{1}{a}\;  (mod\:M)\\

よって、

a^{M-2}

これを求めることで逆元が得られる。
これをPythonで書くと以下のようになる。

pを底とした時の1/aの逆元

a=2
M=7
print(pow(a,M-2,M))

M が素数で，a が M の倍数でない正の整数という条件があるが
M = 1000000007 ,998244353
この条件は M は素数であり、aが 2以上 M 未満の積であるため必ず成り立つといえる。
よってこれは本問題に組み込むことが可能であるといえる。
従ってフェルマーの小定理から逆元を求めるコードが完成した。

＊余談
実は以下のコードでも得られる。

Pythonのみで通るコード

a=2
M=7
print(pow(a,-1,M))

上記のコードでもしっかりと解が得られる。
私の実行環境で試した結果、Pythonにおいて -1のほうが1.7倍くらい早かった
しかし、競プロでよく使われるPyPyでは通らないため、注意が必要

拡張ユークリッドの互除法の理論及び実装

まず、式変形してユークリッドの互除法が適用できる形に変形する。
aを整数、bを任意の整数とする(aは正である必要があるが最終的にmodを取ることで負の数でも問題ない)

a≡\frac{1}{p} \;  (mod\:M)\\
a*p≡1\;  (mod\:M)\\
∴\\
a*p+b*M=1\\

この計算式から、
拡張ユークリッドの互除法はx,yを互いに素の正の整数とするときのa,bを求める手法である。

a*x+b*y=1\\

条件よりx,yを互いに素の正の整数である必要があるが
y=Mであり素数であり、x は 2 以上 M 未満の積であるため本手法が適応可能であるといえる。

この式から、a,bを求めることができれば逆元も求めることができるとなる。
ではコードに変換するための計算過程を示す。
ではおぼろげに浮かんできた67と14で行ってみる。

67=14*4+11\;  => \;67*1+14*(-4)-11=0\;・・・(1)\\

14=11*1+3\;  => \;14*1+11*(-1)-3=0\;・・・(2)\\

11=3*3+2\;  => \;11*1+3*(-3)-2=0\;・・・(3)\\

3=2*1+1\;  => \;3*1+2*(-1)=1\;・・・(4)\\

まず(4)+(3)*(-1)

11*(-1)+3*4=1\;・・・(5)\\

次に(5)+(2)*4

14*4+11*(-5)=1\;・・・(6)\\

最後に(6)+(1)*(-5)

67*(-5)+14*24=1\;・・・(5)\\

よって、a=-5,b=24と求められた。
これを再帰関数を用いて組み上げる。
組むときの流れ

目標とする流れ

extendUclid(67,14)
    extendUclid(14,11)
        extendUclid(11,3)
            extendUclid(3,2)
            return 1, -1
        return -1,4
    return 4, -5
return -5,24

上の図を頼りに再帰関数を組み上げると

def extendUclid(x,y):
    if x%y==1:
        return 1,-(x//y)
    else:
        a,b = extendUclid(y,x%y)
        return b, -(x//y)*b+a
x=67;y=14;
print(extendUclid(x,y))
#(-5, 24)

答えも合い無事完成したことが確認できた。

最後に組み込もう

メモ化もしっかり意識して

フェルマーの小定理を利用した逆元

import math
mod=998244353
n=10000

##関数部分
n1=[1]*(n+1)
for i in range(2,n+1):
    n1[i]=(n1[i-1]*i)%mod
n2=[1]*(n+1)
for i in range(2,n+1):
    n2[i]=pow(n1[i],mod-2,mod)
def comb(N,R,mod=mod):
    return (n1[N]*(n2[N-R]*n2[R])%mod)%mod
##
p=0
for i in range(1,n+1):
    p+=comb(n,i)
    p%=mod
print(p)

拡張ユークリッドによる実装

import math
mod=998244353
n=10000

##関数部分
n1=[1]*(n+1)
for i in range(2,n+1):
    n1[i]=(n1[i-1]*i)%mod

def extendUclid(x,y):
    if x%y==1:
        return 1,-(x//y)
    else:
        a,b = extendUclid(y,x%y)
        return b, -(x//y)*b+a

n2=[1]*(n+1)
for i in range(2,n+1):
    n2[i]=extendUclid(mod,n1[i])[1]
def comb(N,R,mod=mod):
    return (n1[N]*(n2[N-R]*n2[R])%mod))%mod
##
p=0
for i in range(1,n+1):
    p+=comb(n,i)
    p%=mod
print(p)

まとめ

よって、あまりの逆元を用いた組み合わせの総和は上記のようなコードから求めることができた。
初心者用の説明サイトになってしまったが、コンテスト中にこのほしくなって見に来ることがあるかもしれませんね。
以上、お疲れさまでした！！！