Matrix Multiplication

概要

\(N \times N\) の行列 \(A\), \(B\) について、 \(AB = C\) となる \(C\) を計算する。

行列 \(A, B, C\) の \(i\) 行 \(j\) 列の要素をそれぞれ \(a_{ij}, b_{ij}, c_{ij}\) とする時、

\(\displaystyle c_{ij} = \sum_{1 \le k \le N} a_{ik} b_{kj}\)

で計算される。

計算量

\(O(N^3)\)

実装

def matmul(N, A, B, C):
    *TB, = zip(*B) # transpose
    for i in range(N):
        Ai = A[i]
        C[i][:] = (sum(Aik * Bkj for Aik, Bkj in zip(Ai, TBj)) for TBj in TB)
    return C

N = 3
A = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
B = [[1, 0, 1], [2, 1, 0], [0, 0, 1]]
C = [[0]*N for i in range(N)]
matmul(N, A, B, C)
print(C)
# => "[[5, 2, 4], [14, 5, 10], [23, 8, 16]]"

実装 (logical matrix)

from operator import itemgetter
# caluculate a logical matrix C s.t. AB = C
# - A, B and C are N×N logical matrices
def matmul(N, A, B, C):
    F = [itemgetter(j) for j in range(N)]
    for i in range(N):
        B0 = [Bk for Aik, Bk in zip(A[i], B) if Aik]
        C[i][:] = (sum(map(f, B0)) & 1 for f in F)
    return C

N = 3
A = [[1, 0, 1], [0, 1, 1], [1, 1, 0]]
B = [[1, 0, 0], [1, 1, 0], [0, 0, 1]]
C = [[0]*N for i in range(N)]
matmul(N, A, B, C)
print(C)
# => "[[1, 0, 1], [1, 1, 1], [0, 1, 0]]"

Verified

  • AOJ: "2624: Graph Automata Player": source (Python3, 36.71sec)

戻る