행렬 덧셈

import sys


def main():
    input = sys.stdin.readline

    rows, cols = map(int, input().split())

    mat_a = [list(map(int, input().split())) for _ in range(rows)]
    mat_b = [list(map(int, input().split())) for _ in range(rows)]
    mat_c = [[a + b for a, b in zip(row_a, row_b)] for row_a, row_b in zip(mat_a, mat_b)]

    for row in mat_c:
        print(*row)


main()

2차원 배열

Python에서는 list 안에 list를 넣어서 2차원 배열을 표현할 수 있습니다. 흔히 행 방향(위쪽에서 아래쪽으로)을 y축, 열 방향(왼쪽에서 오른쪽으로)을 x축으로 삼고, 각 원소의 위치를 (y, x) 순서쌍으로 표현합니다. 이 관습을 따르면, 아래 코드에서 arr은 2행 3열, brr은 3행 2열인 2차원 배열을 나타냅니다.

arr = [[1, 2, 3], [4, 5, 6]]

brr = [
    [1, 2],
    [3, 4],
    [5, 6],
]

인덱싱은 순서쌍에 정의된 순서대로 하나씩 인덱스를 적용하면 됩니다. 위의 코드에 이어서, 다음과 같이 사용할 수 있습니다.

print(arr[1][0])    # 4
print(brr[-1][-1])  # 6

많지는 않지만, 일부 문제는 방향을 반대로 정의하는 경우가 있으므로 주의해야 합니다.