行列の累乗を繰り返し二乗法で高速計算する方法（Python）

Aru

行列$A$の$n$乗、つまり$A^n$を効率的に計算するプログラムをPythonで実装します。特に、$n$が$10^{18}$のように非常に大きな場合、従来の方法では計算コストが膨大になります。本記事では、繰り返し二乗法を活用し、行列の累乗を高速化するための具体的な手法をわかりやすく解説します。

その他のAtCoderに役立つ記事の一覧は以下にあります。

AtCoderで役立つ記事一覧（Python, Go言語）

Contents

繰り返し二乗法とは
- 基本的なアイデア
どんな時に使う？
フィボナッチ数列の場合の考え方
Pythonで実装
終わりに

繰り返し二乗法とは

繰り返し二乗法は、特に$n$が大きな$A^n$を効率的に計算する効率的に計算するためのアルゴリズムです。この手法を使うと、冪乗の計算を対数的な回数に減少させることができ、計算量を大幅に減らすことができます。

基本的なアイデア

繰り返し二乗法の基本的なアイデアは、指数を２の累乗の和に分割し計算するテクニックです。

例えば、$2^11$の指数11は$8+2+1=11$と表すことができます。

これを使って式を変形すると、以下のようになります。

$$
2^{11} = 2^{8+2+1} = 2^8 \times 2^2 \times 2^1
$$

あらかじめ$2^8$と$2^2$と$2^1$を計算しておけば、この式は3回の乗算だけで計算できます。例えば、$2^4$から$2^8$を計算するのは、$2^8 = 2^4 \times 2^4$なので順に計算していけば計算量はO(logn)となり高速です。

どんな時に使う？

いちばんわかりやすい例がフィボナッチ数列のn番目の値を調べるというやつでしょうか。フィボナッチ数列は、以下のような数列です。

$$
a_1 = 1 , a_2 = 1
$$

$$
a_n = a_{n-1} + a_{n-2}
$$

たとえば、「10,000番目は？」みたいな問いに対してどう答えれば良いでしょうか。$a_3$を計算して、$a_4$を計算して…と順番にやると計算する回数がとても多くなります

これを高速にやるのが行列の累乗になります。

フィボナッチ数列の場合の考え方

いきなり答えに近いですが、例えば、$a_3$, $a_2$を行列で表すと以下になります。

$$
\begin{pmatrix}
a_3\\
a_2
\end{pmatrix} =
\begin{pmatrix}
1&1\\
1&0
\end{pmatrix}
\begin{pmatrix}
a_2\\
a_1
\end{pmatrix}
$$

また、$a_4$, $a_3$は、上の式の結果を用いて計算すると以下になります。

$$
\begin{pmatrix}
a_4\\
a_3
\end{pmatrix} =
\begin{pmatrix}
1&1\\
1&0
\end{pmatrix}
\begin{pmatrix}
a_3\\
a_2
\end{pmatrix}=
\begin{pmatrix}
1&1\\
1&0
\end{pmatrix}
\begin{pmatrix}
1&1\\
1&0
\end{pmatrix}
\begin{pmatrix}
a_2\\
a_1
\end{pmatrix}=
\begin{pmatrix}
1&1\\
1&0
\end{pmatrix}^2
\begin{pmatrix}
a_2\\
a_1
\end{pmatrix}
$$

これをnまでで考えると、以下の式が成り立つことがわかるかと思います。

$$
\begin{pmatrix}
a_{n}\\
a_{n-1}
\end{pmatrix} =
\begin{pmatrix}
1&1\\
1&0
\end{pmatrix}^{n-2}
\begin{pmatrix}
a_2\\
a_1
\end{pmatrix}
$$

ということで、n番目を計算するには、行列の累乗をすれば良いことがわかります。累乗計算は高速に計算する方法がありますので、行列累乗は高速に行うことができます。

Pythonで実装

nが大きくなると桁数が大きくなってしまうので、AtCoderなどでは「xxで割ったあまり」を計算することになります。ということで、ここでは、modで割ったあまりを計算する形式で実装したいと思います。

行列の積

まず、行列の積を計算する関数を用意します。この関数は、特にわかりにくいことはないかと思います。

行列Aと行列Bの乗算を行います。ところで、本来は、行列Aの列数と、行列Bの行数が一致しているかチェックが必要ですが、この関数ではエラー処理を行なっていません。

def matmul(A, B, mod):
    N = len(A)
    K = len(A[0])
    M = len(B[0])

    c = [[0 for _ in range(M)] for _ in range(N)]

    for i in range(N) :
        for j in range(K) :
            for k in range(M) :
                c[i][k] += A[i][j] * B[j][k] 
                c[i][k] %= mod
    return c

累乗

こちらが、メインの行列の累乗を行う関数です。

先ほどの式ですが、先頭に単位行列を置いて考えた方がプログラム的には都合が良いです（n=0回でも問題なく動くので）。ということで、cの初期値を単位行列にしています。

$$
\begin{pmatrix}
1&0\\
0&1
\end{pmatrix}
\begin{pmatrix}
1&1\\
1&0
\end{pmatrix}^n
$$

累乗の計算は、powの計算でよく行われているテクニックを、そのまま行列に拡張した感じです。

def pow_matrix(A, p, mod):
    n = len(A)
    c = [[1 if i == j else 0 for i in range(n)] for j in range(n)]
    while p > 0 :
        if p%2 == 1 :
            c = matmul(c, A, mod)
        A = matmul(A, A, mod)
        p //= 2
    return c

使い方

フィボナッチ数列のn番目を求める場合は、以下のようにして利用します。Aは$N \times N$行列であれば良いので、他にも応用できる関数となっています。活用してください

A = [[1,1],[1,0]]
n = 10
mod = 10**9 + 7
B = pow_matrix(A, n, mod)
print((B[0][0] + B[0][1])%mod

終わりに

行列の累乗は高速化テクニックとして結構有用です（競プロだけでなく、普段でも）

最後に関数全体をつけておきます

def matmul(A, B, mod):
    N = len(A)
    K = len(A[0])
    M = len(B[0])

    c = [[0 for _ in range(M)] for _ in range(N)]

    for i in range(N) :
        for j in range(K) :
            for k in range(M) :
                c[i][k] += A[i][j] * B[j][k] 
                c[i][k] %= mod
    return c

def pow_matrix(A, p, mod):
    n = len(A)
    c = [[1 if i == j else 0 for i in range(n)] for j in range(n)]
    while p > 0 :
        if p%2 == 1 :
            c = matmul(c, A, mod)
        A = matmul(A, A, mod)
        p //= 2
    return c

#python #行列の累乗