フィボナッチ数列を行列を使ってO(log(n))で計算する

(更新 2022/06/24 行列は累乗はよく使うので汎用関数 mtxPowMod(mtx, n, m)を定義してそれを使ったコードを加えました。)

以下のような例題を解くためにフィボナッチ数列を高速で計算する方法はないかと探していたら「行列累乗でフィボナッチ数列を計算する」という投稿が見つかりました。これを参考にしてPythonのプログラムを書いてみました。

例題. フィボナッチ数列を$f(n)$としたとき$f(10^{14})$をmod 100007で求めよ

これをループや再帰関数で求めるのはできないので以下の行列の累乗を使って求めるようにしました。
式はWikipediaのフィボナッチ数にも載っています。

\begin{align}
\large \begin{pmatrix} F_{n+1} & F_n \\\ F_n & F_{n-1} \end{pmatrix}
= \begin{pmatrix} 1 & 1\\\ 1 & 0 \end{pmatrix}^n
\end{align}

従ってnが$2^k$の時は行列の2乗をk回繰り返せばよいので以下のようになります。この時に要注意なのがpythonのnp.arrayの要素はint64なのでMがもっと大きいとオーバフローを起こしてしまいます。arrayを作るときにdtype='object'とすると通常の整数変数と同様に桁数が無制限にできました。

import numpy as np
import math
M = 100007
def fibmtx2n(k):   # n = 2 ** n2
  fm = np.array([[1,1],[1,0]], dtype='object')
  for i in range(k): 
    fm = np.dot(fm,fm) % M
  return fm

このfibmtx2n()を使って任意のnを2のべき乗の和に分解して呼ぶことで計算することが出来ます。

def fibmtx(n):
  fm1 = np.array([[1,0],[0,1]],dtype='object')
  while n > 0:
    l2n = n.bit_length()-1
    fm1 = np.dot(fm1, fibmtx2n(l2n)) % M
    n -= 2**l2n
  return fm1[0][1]

fixbtx()関数を使って以下のように例題の計算を瞬時に行うことが出来ました。

n = 10**14
print(n, fibmtx(n))
# 100000000000000 63949

さらにこれを何度も使うときにはfibmtx2n()関数に@lru_cacheを適用することでさらに早くすることが出来ます。

(* 2021/10/06 コメントを反映してbit_length()に変更)

(以下 2022/06/24に更新)

関数 mtxPowModを使ったコード

pythonのpow関数と同様のmoduloを指定できる行列の累乗関数mtxPowModで、よりシンプルなコードになりました。

import numpy as np
# matrix exponentiation, (mtx:matrix, k: power, M: modular)
def mtxPowMod(mtx,k, M):               
  ret = np.identity(len(mtx), dtype=int)  # identity matrix
  while k > 0:
    if (k % 2) == 1: ret = ret @ mtx % M
    mtx = (mtx @ mtx) % M
    k //= 2
  return ret

def fibn(N, M):  # n-th Fibonacci mod M
  fmtx = np.array([[1,1],[1,0]], dtype=object)
  return mtxPowMod(fmtx,N, M)[0,1]

N, M = 10**14, 100007
print(N, fibn(N, M))
# 100000000000000 63949

（開発環境：Google Colab）

このやり方はProject Euler Problem 304: Primonacci等解くのに役に立ちます