冪乗の計算と二項係数の計算 Try and Error

はじめに

数学的に厳密でない表現があります。

やったこと

• $a^n$ を$10^9+7$で割ったあまりを求める
• 下降階乗冪を$10^9+7$で割ったあまりを求める
• 二項係数を$10^9+7$で割ったあまりを求める

以下、有限体 $F_{p}\ (p=10^9+7)$ 上での演算を考えます。

動機

AtCoderのこれ。
D - Bouquet

非常に大きい数での組み合わせの総数を求める問題です。

\sum_{k=0}^{n}
\biggl(
\begin{matrix}
n \\
k
\end{matrix}
\biggl)
\ =\ 
(1+1)^n \ =\ 2^n
\ \ \ \ \ \ (0\leq n\leq 10^9)

問題文を読むとわかりますが、ここからさらに

\biggl(
\begin{matrix}
n \\
a
\end{matrix}
\biggl)_{,}\ \ \ 
\biggl(
\begin{matrix}
n \\
b
\end{matrix}
\biggl)
\ \ \ \ \ (1≤a<b≤min(n,2×10^5))

を引きます。なので、以下2点が必要になります。

• 非常に大きい数(10^9)の冪乗を求められる
• まあまあ大きい数(?)(2×10^5)の二項係数を求められる
  

愚直な方法

素直に2^(10^9)を求める

Prelude> snd $ divMod (2^1000000000) (10^9 + 7)
140625001

先生「はい、みんなが計算結果を出すまで12秒くらいかかりました。」
余裕のTLEです。

どうにか二項係数を求めてみる

• 二項係数を計算する関数を再帰的に定義してみる

combination :: (Integral a) => a -> a -> a
combination 0 _ = 1
combination _ 0 = 1
combination m n
    | m == n = 1
    | m > n  = snd $ (+) (snd $ combination (m-1) n `divMod` p) (snd $ combination (m-1) (n-1) `divMod` p) `divMod` p
    where p = 10^9 + 7

・・・書いてて絶対ダメと思いつつ念のため一応動作確認しましたが$_{30}C_{6}$とかでもう怪しかったです。。

• 下降階乗冪を計算する関数を作ってみる

-- fact' n k でn*(n-1)..*(n-k+1) を計算　combinationの分子及び分母を計算
fact' :: (Integral a) => a -> a -> a
fact' n 1 = m
fact' n k
    | n < k  = -1
    | n >= k = snd $ (n * fact' (n-1) (k-1)) `divMod` p
    where p = 10^9 + 7

これはぱぱっと動いてくれました。

*Lib> fact' 1000000000 200000
555005661

が、割り算しなくてはいけないわけで、どうしよう・・

• 逆元を求める関数を作ってみる

$a$で割るということは$a^{-1}$を掛けるということです。
また、Fermatの小定理より逆元は$a^{p-2}$で与えられます。
-> つまり冪乗が高速に求められない限り逆元も求まりません！終わった・・・・・

※Euclidの互除法で求められるようですが今回は見送りました。

改良した方法

冪乗

2分割しながら計算していきます。
例えば、$2^{11}$を次のように計算していきます。

$2^{11} = 2^{5} \times 2^{5} \times 2$
$2^{5} = 2^{2} \times 2^{2} \times 2$
$2^{2} = 2^{1} \times 2^{1}$

これなら計算量が$O(\log n)$になる気が・・？実装してみます。

power' :: (Integral a) => a -> a -> a
power' 1 _ = 1
power' a 1 = a
power' a b = let ret = snd $ divMod (x * x) p
             in if even b then ret else snd $ divMod (ret * a) p
             where p = 10^9 + 7
                   x = power' a $ fst (b `divMod` 2)

*Lib> power' 2 1000000000
140625001

めちゃめちゃ早くなった！！！おもしろーい

そのほかの関数

次はpower'を使って逆元を求める関数を実装します。

inv :: (Integral a) => a -> a -> a
inv n p = if n <= 0 || p <= 1
              then -1
              else power' n (p-2)

*Lib> inv 3746284 1000000007
319182691

早い！OK！最後にinvを使って二項係数を計算する関数を実装します。

combination' :: (Integral a) => a -> a -> a
combination' 0 _ = 1
combination' _ 0 = 1
combination' m 1 = m
combination' m n
    | m < n  = 0
    | m == n = 1
    -- 1000C999 を 1000C1 にした方が嬉しいよね
    | n > threshold  = combination' m (m-n)
    | n <= threshold = snd $ ((fact' m n) * (inv (fact' n n) divisor)) `divMod` divisor
    where threshold = fst $ divMod m 2
          divisor = 10^9 + 7

*Lib> combination 1000000000 200000
178422379

$_{30}C_{6}$で悲鳴を上げていたcombinationくんはもういません。

おわりに

アルゴリズム周り、及びHaskellに関してド素人もいいところなのですが、
自分なりに試行錯誤して高速化をしていく作業はとても楽しいですね。
Haskellもとても楽しいです。学び始めて数週間とかでまだまだすぎますが挫折しないように頑張ろう・・！