プログラミング
記事内に商品プロモーションを含む場合があります

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

Aru

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

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

繰り返し二乗法とは

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

基本的なアイデア

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

例えば、$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

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

ABOUT ME
ある/Aru
ある/Aru
IT&機械学習エンジニア/ファイナンシャルプランナー(CFP®)
専門分野は並列処理・画像処理・機械学習・ディープラーニング。プログラミング言語はC, C++, Go, Pythonを中心として色々利用。現在は、Kaggle, 競プロなどをしながら悠々自適に活動中
記事URLをコピーしました