当前位置：首页 > article >正文

【算法】矩阵乘法与矩阵快速幂 python

article 2025/3/24 15:04:28

矩阵乘法

模板

https://www.acwing.com/problem/content/1304/

import sys  
input = lambda: sys.stdin.readline().strip()  
  
n, m = map(int, input().split())  
A = [list(map(int, input().split())) for _ in range(n)]  
p = int(input())  
B = [list(map(int, input().split())) for _ in range(m)]  
  
ans = [[0] * p for _ in range(n)]  
for i in range(n):  
    for j in range(p):  
        for k in range(m):  
            ans[i][j] += A[i][k] * B[k][j]  
              
for row in ans:  
    print(*row)

矩阵快速幂

快速幂算法的核心是利用指数的二进制分解。

def mat_pow(A, k):  
    n = len(A)  
    ans = [[1 if i == j else 0 for j in range(n)] for i in range(n)]  
    while k > 0:  
        if k & 1:  
            ans = mat_mul(ans, A)  
        A = mat_mul(A, A)  
        k >>= 1  
    return ans

模板

https://www.luogu.com.cn/problem/P3390

import sys  
input = lambda: sys.stdin.readline().strip()  
  
mod = 10 ** 9 + 7  
def mat_mul(A, B):  
    n = len(A)  
    m = len(A[0])  # m = len(B)  
    p = len(B[0])  
  
    C = [[0] * p for _ in range(n)]  
    for i in range(n):  
        for j in range(p):  
            for k in range(m):  
                C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % mod  
    return C  
  
def mat_pow(A, k):  
    n = len(A)  
    ans = [[1 if i == j else 0 for j in range(n)] for i in range(n)]  
    while k > 0:  
        if k & 1:  
            ans = mat_mul(ans, A)  
        A = mat_mul(A, A)  
        k >>= 1  
    return ans  
  
n, k = map(int, input().split())  
A = [list(map(int, input().split())) for _ in range(n)]  
ans = mat_pow(A, k)  
for row in ans:  
    print(*row)