[백준][파이썬] 28438번. 행렬 연산 (행렬 계산하기)

문제

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

 

풀이

첫 번째 시도

한 줄씩 입력을 받으면서 '1 r v'인지 '2 c v'인지 구분하며 행렬 연산을 수행했다.

N, M, Q = map(int, input().split())
mat = [[0 for _ in range(M)] for _ in range(N)]

for _ in range(Q) :
  i, j, k = map(int, input().split())
  if i == 1 :
    r, v = j, k
    for l in range(len(mat[r-1])) :
      mat[r-1][l] += v
  else :
    c, v = j, k
    for l in range(len(mat)) :
      mat[l][c-1] += v

for i in mat :
  print(*i)

→ 시간 초과

두 번째 시도 (성공)

입력을 받을때마다 for문을 사용해서 행렬 연산을 진행하는 부분때문에 시간 초과가 발생했다.

N x M 모양의 행렬을 만든 후 계산하는게 아니라, row와 col로 쪼개서 생각해주었다.

row[i] + col[j]는 mat[i][j]의 값과 동일하다.

import sys
input = sys.stdin.readline

N, M, Q = map(int, input().split())
row = [0 for _ in range(N)]
col = [0 for _ in range(M)]

for _ in range(Q) :
  i, j, k = map(int, input().split())
  if i == 1 :
    r, v = j-1, k
    row[r] += v
  else :
    c, v = j-1, k
    col[c] += v

for i in range(N) :
  for j in range(M) :
    print(row[i] + col[j], end = ' ')
  print()