matrix.py

from __future__ import annotations
import operator

MatType = list[list[float]]

# N matmul
# I {"version": "1.1", "import":["operator"], "const": ["MatType"]}
def matmul(mat_a: Matrix, mat_b: Matrix, mod: int = 0) -> Matrix:
    if mod == 0:
        return [
            [sum(map(operator.mul, a_row, b_col)) for b_col in zip(*mat_b)]
            for a_row in mat_a
        ]
    return [
        [sum(map(operator.mul, a_row, b_col)) % mod for b_col in zip(*mat_b)]
        for a_row in mat_a
    ]

• 행렬의 곱셈을 구하는 함수. 나이브한 O(n^3)알고리즘을 사용
• map을 써서 계산하는 것이 컴프리헨션을 써서 sum(a*b for a, b in zip(a_row, b_col)) 처럼 쓰는것보다 빠르다

# N matpow
# I {"version": "1.1", "const": ["MatType"], "func": ["matmul"]}
def matpow(mat: Matrix, n: int, mod: int = 0) -> MatType:
    res = [[0] * len(mat) for _ in mat]
    for i, row in enumerate(res):
        row[i] = 1
    m = [row[:] for row in mat]
    for c in bin(n)[:1:-1]:
        if c == '1':
            res = matmul(res, m, mod)
        m = matmul(m, m, mod)
    return res

• 행렬의 거듭제곱을 구하는 함수. 바이너리 익스포넨션을 사용해서 O(logn)번의 행렬곱셈을 수행한다.

# N Matrix
# I {"version": "0.2", "future": ["annotations"], "func": ["identity", "matmul", "matpow"]}
class Matrix:
    """A simple wapper class of matrix operation functions."""

    __slots__ = ('mat',)
    _mod = 0

    @classmethod
    def set_mod(cls, mod: int):
        cls._mod = mod

    def __init__(self, mat_or_size):
        if isinstance(mat_or_size, int):
            self.mat = [[0] * mat_or_size for _ in range(mat_or_size)]
        else:
            self.mat = [r[:] for r in mat_or_size]

    @classmethod
    def identity(cls, size: int):
        mat = [[0] * size for _ in range(size)]
        for i, row in enumerate(mat):
            row[i] = 1
        return cls(mat)

    def __getitem__(self, row: int):
        return self.mat[row]

    def __matmul__(self, other: Matrix):
        return Matrix(matmul(self.mat, other.mat, self._mod))

    def __pow__(self, n, mod=0):
        return Matrix(matpow(self.mat, n, (mod or self._mod)))

    def __mod__(self, mod):
        return Matrix([[item % mod for item in row] for row in self.mat])

• @과 **연산자 또는 pow()로 곱과 거듭제곱을 계산할 수 있게 만든 행렬 클래스. matmul과 matpow를 랩핑하도록 구현되어있다.
• matmul이나 matpow로 코드를 작성하는 것이 가독성이 떨어질 경우에 사용한다.
• 행렬곱 이후에 모듈러를 취해줘야 할 경우, 명시적으로 매번 @ 연산을 해주는 방법과 Matrix.set_mod()를 통해 전역적인 모듈러 값을 세팅해주는 방법이 있다. 대부분의 경우, 후자가 간단하다.
• 리스트의 리스트 형태로 실제 행렬을 저장하는 mat변수는 의도적으로 public으로 설정되었다.