フェルマーの小定理を使って mod の下での組み合わせ総数を求めるアルゴリズムを python で書いてみる

{}_nC_r素数 p で割った余りを求める方法について解説します. {}_nC_r は"n 個の異なるものからr 個を選ぶときの組み合わせの数"を表しており, 階乗記号 k!=k\times(k-1)\times \cdots \times 2\times 1を用いて,

{}_nC_r=\dfrac{n!}{(n-r)! r!}

と表されます. 階乗の大きさは指数関数的に増大していくので, 分母と分子は非常に大きな数になりえます. 素数 p で割った余りさえわかればよいことを用いて, どうにか計算量を減らしていきましょう.

前回記事では法の下での掛け算について
 (a \times b) \textrm{ mod } p = ( (a \textrm{ mod } p ) \times (b \textrm{ mod } p)) \textrm{ mod } p
という性質があることを紹介しました. 一方で法の下での割り算では, \frac{a}{b}が整数だと期待できたとしてもこのような性質は成り立ちません. しかし, b の法 p における逆元 b^{-1}, つまり, 整数であって(b\times b^{-1}) \textrm{ mod } p =1をみたすものを求めることができたならば,

 \begin{split} \dfrac{a}{b} \textrm{ mod } p &=(a \times \dfrac{1}{b}) \textrm{ mod } p \\
&= (a \textrm{ mod } p) \times (\dfrac{1}{b} \times (b\times b^{-1}) \textrm{ mod } p) \textrm{ mod } p \\
&= ( (a \textrm{ mod } p) \times (b^{-1} \textrm{ mod } p) ) \textrm{ mod } p\end{split}

によって求めることができます. 実は p素数であるという性質と前回の指数計算アルゴリズムから, b^{-1}を 効率よく求めることができます.

フェルマーの小定理

フェルマーの小定理について, 詳しくはフェルマーの小定理 - Wikipediaをご覧ください. フェルマーの小定理から特に,

p素数, 自然数  b1\le b < p をみたすとき,

 (b^{p-2}\times b) \textrm{ mod } p =1

つまり,  b^{-1}= b^{p-2} \textrm{ mod } p が成立します. 前回の指数計算のアルゴリズムでは p が大きくても高々  O(\log{p}) の計算量ですみますので, 十分に高速です.


そこで, 次のようなアルゴリズム{}_nC_k \textrm{ mod } p を求めます:

  1. n!, k!, (n-k)! を求める
  2.  n! \textrm{ mod } p を求める.
  3. (k!)^{p-2}, ((n-k)!)^{p-2} \textrm{ mod } p を mod の下での指数計算によって求める.
  4. それぞれを掛けて, p で割った余りを求める.

以下がpythonのコードになります.

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

ただし p は大きな素数で, 特に n より大きいものを選んでください. さもなくば n!p の倍数になり, フェルマーの小定理を用いることが出来ません.