[프로그래머스] 👩‍🏫 행렬의 곱셈 / Python

👩‍🏫 행렬의 곱셈 문제 풀어보기

😃 나의 코드

def solution(arr1, arr2):
    answer = []
    for idx1 in range(len(arr1)):
        row = []
        for idx2 in range(len(arr2[0])):
            tmp = 0
            for idx3 in range(len(arr1[0])):
                tmp += arr1[idx1][idx3] * arr2[idx3][idx2]
            row.append(tmp)
        answer.append(row)
    return answer

제한 조건에서 곱할 수 있는 배열만 주어진다고 했으니, 두 배열을 A X B, B X C로 해보겠습니다. 그리고 두 행렬의 곱으로 나온 배열의 크기는 A X C가 될 것입니다. 결과 배열의 왼쪽 위부터 오른쪽 아래까지 만들어 나가면 되는 쉬운 문제입니다.

def solution(arr1, arr2):
    A = len(arr1)
    B = len(arr1[0])
    C = len(arr2[0])
    return [[sum([arr1[idx1][idx3] * arr2[idx3][idx2] for idx3 in range(B)]) for idx2 in range(C)] for idx1 in range(A)]

제가 만든 코드정도로 짜도 돌아가긴 하지만, for loop를 이용하여 list를 좀 더 쉽게 만들어주는 방법인 Python의 list comprehension을 이용해서 좀 더 간단하게 만들어보겠습니다.

def productMatrix(A, B):
    return [[sum(a*b for a, b in zip(A_row,B_col)) for B_col in zip(*B)] for A_row in A]

하지만 더 간단하게 만들 수도 있습니다. arr2 열에 대한 배열을 zip(*B)를 이용하면 코드가 간단해집니다.

Written with StackEdit.