行列と行列の積

matmu.py

from typing import List, Union
import unittest
from operator import mul

Matrix = List[List[Union[int, float]]]


def matmul(mat1: Matrix, mat2: Matrix) -> Matrix:
    return [[sum(map(mul, row, column)) for column in zip(*mat2)] for row in mat1]


class TestMatmul(unittest.TestCase):
    def test_正方行列(self):
        mat1 = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
        mat2 = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
        expected = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
        self.assertEqual(matmul(mat1, mat2), expected)

    def test_２つの行列の積が計算できる(self):
        mat1 = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 0, 1, 2], [3, 4, 5, 6]]
        mat2 = [[5, 3, 6, 7], [4, 2, 8, 1], [8, 0, 1, 4], [7, 2, 9, 2]]
        expected = [
            [65, 15, 61, 29],
            [161, 43, 157, 85],
            [67, 31, 73, 71],
            [113, 29, 109, 57],
        ]
        self.assertEqual(matmul(mat1, mat2), expected)


if __name__ == "__main__":
    unittest.main()