Pythonで行列の累乗を実装する

tadanori

行列$A$の$n$乗の$A^n$を計算するプログラムです。nが$10^18$などになるとかなり重い処理になりますので工夫が必要となります。

ここでは、実装方法について解説します

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

どういうときに使うのか？

いちばんわかりやすい例がフィボナッチ数列の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 #行列の累乗