[백준] #21610 마법사 상어와 비바라기 (python)

내가 푼 코드에서는 시간 초과가 나서 정답 코드를 찾아보았다

 

틀린 코드

N, M = list(map(int, input().split()))
rain = []
for i in range(N):
    tmp = list(map(int, input().split()))
    rain.append(tmp)

dir = {1:[0,-1], 2:[-1,-1], 3:[-1,0], 4:[-1,1], 5:[0,1], 6:[1,1], 7:[1,0], 8:[1,-1]}
cloud = [[N-1,0],[N-1,1],[N-2,0],[N-2,1]]

for _ in range(M):
    d, s = map(int, input().split())
    for c in cloud:
        c[0], c[1] = (c[0] + dir[d][0]*s) % N, (c[1] + dir[d][1]*s) % N
        rain[c[0]][c[1]] += 1

    # step 3
    nxt_dir = [2, 4, 6, 8]
    for c in cloud:
        for nxt in nxt_dir:
            nxt_c_0 = c[0] + dir[nxt][0]
            nxt_c_1 = c[1] + dir[nxt][1]
            if 0 <= nxt_c_0 < N and 0 <= nxt_c_1 < N:
                if rain[nxt_c_0][nxt_c_1] >= 1:
                    rain[c[0]][c[1]] += 1
    
    # step 4
    new_cloud = []
    for i in range(N):
        for j in range(N):
            if rain[i][j] >= 2 and [i, j] not in cloud:
                rain[i][j] -= 2
                new_cloud.append([i, j])

    cloud = new_cloud

print(sum(map(sum,rain)))

 

타임아웃이 발생한 이유는 step4에서 `if [i, j] not in cloud` 때문이었다..

if로 리스트를 탐색할 때 time complexity가 O(N)이라는걸 몰랐었다... (갈길이 멀다)

그래서 이 부분을 visited 배열을 선언하여 처리해주니 해결되었다.

 

정답 코드

N, M = list(map(int, input().split()))
rain = []
for i in range(N):
    tmp = list(map(int, input().split()))
    rain.append(tmp)

dir = {1:[0,-1], 2:[-1,-1], 3:[-1,0], 4:[-1,1], 5:[0,1], 6:[1,1], 7:[1,0], 8:[1,-1]}
cloud = [[N-1,0],[N-2,0],[N-1,1],[N-2,1]] 

for _ in range(M):
    d, s = map(int, input().split())
    visited = [[0]*N for _ in range(N)]
    new_cloud = []
    
    while cloud :
        x, y = cloud.pop() 
        nx, ny = (x + dir[d][0]*s) % N, (y + dir[d][1]*s) % N 
        rain[nx][ny] += 1
        visited[nx][ny] = 1
        new_cloud.append([nx,ny])

    nxt_dir = [2, 4, 6, 8]
    for nx, ny in new_cloud:   
        for nxt in nxt_dir:
            mx, my = nx + dir[nxt][0], ny + dir[nxt][1]
            if 0 <= mx < N and 0 <= my < N and rain[mx][my]: 
                rain[nx][ny] += 1

    for i in range (N): 
        for j in range (N):
            if rain[i][j] >= 2 and not visited[i][j]:
                rain[i][j] -= 2
                cloud.append([i,j])

print(sum(map(sum,rain)))

 

구현할 때 time complexity를 잘 고려해야한다는걸 다시한번 깨닫게 해준 문제,,

그것만 아니라면 문제 조건 보고 꼼꼼하게만 구현하면 크게 어려운 점은 없었다.