취알스 9주차 기타 알고리즘 - 5/5

5549 행성 탐사

import sys
input = sys.stdin.readline

n, m = map(int, input().split())
k = int(input())
arr=[]
for i in range(n):
    arr.append(list(input()))
pp = [[[0,0,0] for i in range(m+1)] for j in range(n+1)] 
# 2차원 prefix sum 문제에서는 행, 열이 한 줄씩 더 필요함!! 
for i in range(n):
    for j in range(m):
        for l in range(3):
            pp[i+1][j+1][l] = pp[i+1][j][l] +pp[i][j+1][l]- pp[i][j][l]
        if arr[i][j]=='J':
            pp[i+1][j+1][0] += 1
        elif arr[i][j]=='O':
            pp[i+1][j+1][1] +=1
        elif arr[i][j]=='I':
            pp[i+1][j+1][2] += 1
for _ in range(k):
    a, b, c, d = map(int, input().split())
    answer = [0,0,0]
    for i in range(3):
        answer[i] = pp[c][d][i]-pp[a-1][d][i] - pp[c][b-1][i] + pp[a-1][b-1][i]
    print(answer[0], answer[1], answer[2])

주의

for i in range(n):
    for j in range(m):
        planet[i][j+1] = planet[i][j] # 여기
        if arr[i][j]=='J':
            planet[i][j+1][0] = planet[i][j][0]+1
        elif arr[i][j]=='O':
            planet[i][j+1][1] = planet[i][j][1]+1
        elif arr[i][j]=='I':
            planet[i][j+1][2] = planet[i][j][2]+1
        print(planet)

주석으로 표시해놓은 곳을 보면 값이 들어가지 않고 주소값이 들어가는지 값이 변경될 때 원래 배열도 변경됨.

풀이

일단 당연히 배열의 모든 원소를 다 보면서 세면 O(KNM)의 시간이 걸려서 시간초과가 발생함. 그래서 행마다 prefix sum을 저장하여 풀어보았다. 이론상 O(KN)의 시간이 걸릴텐데 시간초과가 났다. 결국 답을 보고 말았다..

우리가 구하는 부분은 4번이고, 전체에서 (1+3) + (1+2) - 1을 해주면 된다. dp라고 생각할 수도 있는듯?? 어쨌든 2차원 prefix sum 문제를 풀 수 있게 되었다.