行列累乗

• mat_mul(A, B, mod): 行列の積 $AB$ を mod で求める。 $A$ が $N \times M$、 $B$ が $M \times K$ のとき $O(NMK)$
• mat_pow(A, k, mod): 正方行列 $A$ の $k$ 乗（ $k \ge 0$）を求める。 $A$ が $K \times K$ のとき $O(K^3 \log k)$
• mod の既定値は $998244353$

$K$ 項間の線形漸化式は、 $K \times K$ の行列の累乗で第 $N$ 項を $O(K^3 \log N)$ で求められます。

例（フィボナッチ数列の第 $n$ 項。 $F_0 = 0$, $F_1 = 1$）:

A = [[1, 1], [1, 0]]
f = mat_pow(A, n, 10**9 + 7)[0][1]

コード

def mat_mul(A, B, mod=998244353):
    Bt = list(zip(*B))
    C = []
    for row in A:
        C.append([sum(a * b for a, b in zip(row, col)) % mod for col in Bt])
    return C


def mat_pow(A, k, mod=998244353):
    n = len(A)
    R = [[int(i == j) for j in range(n)] for i in range(n)]
    while k:
        if k & 1:
            R = mat_mul(R, A, mod)
        A = mat_mul(A, A, mod)
        k >>= 1
    return R