行列累乗
(NumberTheory/misc/matrix_pow.py)

• View this file on GitHub
• Last update: 2021-06-07 19:35:09+09:00

概要

行列の累乗を求めるアルゴリズム。

使い方

matrix_pow(A: Sequence[Sequence[int]], k: int, MOD: int) -> List[List[int]]
$n$ 次正方行列 $A$ について $A^k$ を返す。行列の各要素は $\mathrm{MOD}$ で割った余りをとる。計算量 $O(n^3 \log k)$

matvec_mul(A: Sequence[Sequence[int]], vec: Sequence[int], MOD: int) -> List[int]
サイズ $n \times m$ の行列 $A$ と $m$ 次元のベクトル $v$ (vec) について $Av$ を返す。計算量 $O(nm)$

Verified with

• TestCase/unittest/matrix_pow.unittest.test.py

Code

def matrix_pow(A, k, MOD):
    def matrix_mul(A, B, MOD):
        C = [[0] * size for _ in range(size)]
        for i in range(size):
            for k in range(size):
                for j in range(size):
                    C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % MOD
        return C

    if len(A) != len(A[0]):
        raise ValueError('square matrix expected')

    size = len(A)
    B = [[0] * size for _ in range(size)]
    for i in range(size):
        B[i][i] = 1
    while k > 0:
        if k & 1:
            B = matrix_mul(A, B, MOD)
        A = matrix_mul(A, A, MOD)
        k = k // 2
    return B


def matvec_mul(A, vec, MOD):
    if len(A[0]) != len(vec):
        raise ValueError('dimension mismatch between matrix and vector')

    h, w = len(A), len(A[0])
    res_vec = [0] * h
    for i in range(h):
        for j in range(w):
            res_vec[i] += A[i][j] * vec[j]
            res_vec[i] %= MOD
    return res_vec

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.12.4/x64/lib/python3.12/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir]}).decode()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.12.4/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/python.py", line 96, in bundle
    raise NotImplementedError
NotImplementedError

Back to top page