백준 2740: 행렬 곱셈 [python]

https://www.acmicpc.net/problem/2740

문제 접근

행렬 두개를 입력받아, 행렬의 곱을 연산하여 출력하는 문제이다.

행렬의 곱은 A 행렬의 열 개수 = B 행렬의 행 개수여야 곱셈이 가능하다.

예시 1의 경우
1 2
3 4
5 6
3행 2열로 이루어져있는 A와

-1 -2 0
0 0 3
2행 3열로 이루어져 있는 B

A행렬의 열의 개수 = 2, B 행렬의 행 개수 = 2으로 같기에 곱셈이 가능하다.

A x B 행렬 C의 크기는 A 행렬의 행 x B 행렬의 열 인 3x3 으로 계산된다.

해답 및 풀이

import sys
input = sys.stdin.readline

# 1. A 행렬 입력 받기
N, M = map(int, input().split())  # A는 N행 M열
A = [list(map(int, input().split())) for _ in range(N)]

# 2. B 행렬 입력 받기
M2, K = map(int, input().split())  # B는 M2행 K열
B = [list(map(int, input().split())) for _ in range(M2)]

# 3. 결과 행렬 C (N x K)를 0으로 초기화
C = [[0] * K for _ in range(N)]

# 4. 행렬 곱 연산 수행
for i in range(N):         # A의 i행
    for j in range(K):     # B의 j열
        for k in range(M): # A의 열 = B의 행
            C[i][j] += A[i][k] * B[k][j]

# 5. 결과 출력
for row in C:
    print(*row)