프로그래머스 | 행렬의 곱셈 [파이썬 python]

프로그래머스

 
행렬곱에 대한 이해가 있어야 문제를 해결할 수 있으므로, 행렬곱에 대한 자세한 내용은 여기를 참고하길 바란다.

• 두 행렬 arr1, arr2의 크기가 각각 m×n, n×r 일 때, arr1의 열과 arr2의 행의 크기가 같을 때만 행렬의 곱셈이 가능하다.
• arr1과 arr2 곱셈 결과 나오는 행렬의 크기는 m×r이다.

 
코드 - 방법 1)

def solution(arr1, arr2):
    answer = []
    m, n, r = len(arr1), len(arr1[0]), len(arr2[0])

    for i in range(m):
        arr = arr1[i]
        result = []
        for j in range(r):
            hap = 0
            for k in range(n):
                hap += arr[k] * arr2[k][j]
            result.append(hap)
        answer.append(result)

    return answer

행렬곱의 과정을 하나하나 작성한 코드이다.
arr1에서 하나의 행을 가져와 arr로 저장한 후, arr2의 하나의 열과 곱해 hap을 구하고 이 hap들을 하나의 행에 대한 곱셈 결과인 result에 추가한다.

위 그림에서 10+6+6이 hap이며, 한 행에 대한 hap들을 저장한 리스트 [22, 22, 11]이 result이다.
 
코드 - 방법 2) answer 크기 지정

def solution(arr1, arr2):
    m, n, r = len(arr1), len(arr1[0]), len(arr2[0])
    answer = [[0 for _ in range(r)] for _ in range(m)] # m * r 크기의 행렬

    for i in range(m):
        for k in range(r):
            for j in range(n):
                answer[i][k] += arr1[i][j] * arr2[j][k]
    
    return answer

두 행렬 m×n, n×r 곱셈 결과는 m×r 이 되는 것을 알기 때문에 answer의 크기를 먼저 정의할 수 있다.
answer의 해당 위치에 arr1과 arr2의 요소를 곱해 누적해 주면 좀 더 깔끔하게 코드를 작성할 수 있다.
 
다른 사람의 코드

def solution(arr1, arr2):
    return [[sum(i*j for i, j in zip(row, col)) for col in zip(*arr2)] for row in arr1]

먼저 *arr2와 zip(*arr2)를 출력해 보면 다음과 같다.

*arr2, zip(*arr2)

리스트 앞에 *을 붙이면 unpacking을 할 수 있다. 쉽게 말해 대괄호([]) 하나를 벗겨내는 것이다.
arr2=[[5, 4, 3], [2, 4, 1], [3, 1, 1]]는 2차원 배열이었으므로 *arr2을 하면  [5, 4, 3] [2, 4, 1] [3, 1, 1]이 된다.
*arr2를 zip을 사용해 리스트들을 묶어주면, (5, 2, 3) (4, 4, 1) (3, 1, 1)로 같은 위치에 있는 요소들끼리 묶이게 된다.
(5, 2, 3) (4, 4, 1) (3, 1, 1)은 arr2에서 각 열에 해당하는 요소이므로 arr1의 행과 곱해주면 우리가 원하는 결과를 구할 수 있다.
이때도 zip을 사용해 row와 col을 묶어 하나씩 요소를 불러와 곱해주었다.