PSM

Pプログラミング S初心者の Mメモ書き

Python3で組み合わせ計算を高速化

初めに

Pythonで組み合わせ計算 {}_n\mathrm{C}_rを計算しようとすると,少々問題が発生します.
nやrが小さいうちはいいんですが,これが大きくなってくると計算が非常に長くなってくるのです.
実務や研究では少々計算が遅くても別に問題はないと思うのですが,競技プログラミングだとそれが命取りになってくるんですね……

組み合わせ計算のための元(と逆元)のテーブルの計算

階乗をメモしておく

というわけで,計算が重いなら予め計算した結果をテーブルにしといちゃいましょうというテクニックがあるそうです.

さて
 {}_n\mathrm{C}_r = \frac{n!}{r!(n-r)!}
からわかるように,単純に考えれば元 x! 1 \le x \le N(Nは考えられる最大)で計算してリストに書き込んでおけばよいことになります.
よって

#元
N  = 10
g =  [math.factorial(i) for i in range(N) ]

のように求めておけば,例えばあとで10!の値が欲しければg[10]を呼び出すことにより,O(1)で組み合わせ計算を行うことができます.

上のテーブルを使えば組み合わせ計算は

def combination(n,r):
    return int(g[n]/(g[r]*g[n-r]))

のように求めることができます.


出力できる数値に最大値の制限がある場合(競プロ特有の話)

上のようにテーブルに記憶させた場合でも,nやrが大きくなってくると,そもそもの階乗の計算量が膨大になります.

ここで,出力に制限がある状態にはさらに高速にテーブルを作成することができるようです.
具体的には「出力は 10^9+7で割った余りとせよ」のような場合のことです.


 {}_n\mathrm{C}_r = \frac{n!}{r!(n-r)!}
について,元 x!と逆元 (x!)^{-1} 1 \le x \le N(Nは考えられる最大)で計算しておくことを考えます.

 10^9+7素数なので,フェルマーの小定理より
 x ≡ x^{10^9 + 7} \ \ (\mathrm{mod} \ 10^9 + 7)
両辺を x^2で割って
 x^{-1} ≡ x^{10^9 + 5} \ \ (\mathrm{mod} \ 10^9 + 7)
となっています.

元と逆元のテーブルは,以下のように作成し,使うことができるようです.
(下記のコードの元は自分で考えたものではありません)

mod = 10**9+7  #出力の制限
N = 10**4 
g1 = [1, 1]  # 元テーブル
g2 = [1, 1]  #逆元テーブル
inverse = [0, 1]  #逆元テーブル計算用テーブル

for i in range( 2, N + 1 ):
    g1.append( ( g1[-1] * i ) % mod )
    inverse.append( ( -inverse[mod % i] * (mod//i) ) % mod )
    g2.append( (g2[-1] * inverse[-1]) % mod )

#以下のcom関数を用いて組み合わせを計算
def com(n, r, mod):
    if ( r<0 or r>n ):
        return 0
    r = min(r, n-r)
    return g1[n] * g2[r] * g2[n-r] % mod

関連する定理など

二分累乗法

xをn乗したいとき,普通に計算するとn回累乗計算を行うことになるので,計算量は \mathrm{O}(n)になります.
これを高速にする二分累乗法というものがあり,Python3で以下のように実装できます.

# 二分累乗法
def power(x, n):
    y=1
    while( n>0 ):
        if ( n%2==0 ):
            x = x * x
            n = n/2
        else:
            y = y * x
            n = n-1
    return y

nを二進法展開するように2でどんどん割っています.
1.奇数のときは答えにxを掛けてnから1引く
2.偶数のときは答えはそのままにxを2乗してnを1/2にする
という操作を繰り返すのですが,最低でも2回に1回はnが1/2になるので計算量が  \mathrm{O}(\log_2 n)に減少するというわけです.

より数学的な(でも簡略な)説明は以下.

nを二進法で展開すると
 n  = 2^{s_1} + 2^{s_2} + \cdots + 2^{s_k}  ただし   0 \le s_1 < s_2 < \cdots < s_k
となる.これより x^n
 \begin{split} x^n  &= x^{2^{s_1} + 2^{s_2} + \cdots + 2^{s_k}} \\ &= x^{2^{s_1}}x^{2^{s_2}} \cdots x^{2^{s_k}} \end{split}
となる.このとき,各々の x^{2^{s_i}} x^{2^{s_{i-1}}}から順に計算していけるので,計算量は  \mathrm{O}(\log_2 n)になる.

フェルマーの小定理

wikipedia:フェルマーの小定理

p を素数とし、a を p の倍数でない整数(a と p は互いに素)とするときに、
 a^{p-1} ≡ 1 \ \ (\mathrm{mod} \ p)
すなわち、a の p − 1 乗を p で割った余りは 1 であるというもの