[파이썬, Python] 백준 21610: 마법사 상어와 비바라기

문제

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

 

코드

import sys
input = sys.stdin.readline

n, m = map(int, input().split())
buckets = [list(map(int, input().split())) for _ in range(n)]

cloud = [[n - 1, 0], [n - 1, 1], [n - 2, 0], [n - 2, 1]]

def move_cloud():
    move_cloud = []
    order = [(0, 0), (0, -1), (-1, -1), (-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1)]

    d, s = map(int, input().split())

    dir = order[d]

    for i in range(len(cloud)):
        new_i0 = (cloud[i][0] + dir[0] * s) % n
        new_i1 = (cloud[i][1] + dir[1] * s) % n

        move_cloud.append([new_i0, new_i1])
        visited[new_i0][new_i1] = True

    return move_cloud

def rain():
    global cloud

    for i in range(len(cloud)):
        buckets[cloud[i][0]][cloud[i][1]] += 1

        del_cloud.append([cloud[i][0], cloud[i][1]])

    cloud = list()

def water_copy():
    for i in del_cloud:
        r, c = i

        dr = [1, 1, -1, -1]
        dc = [1, -1, 1, -1]

        for i in range(4):
            nr = r + dr[i]
            nc = c + dc[i]

            if 0 <= nr < n and 0 <= nc < n and buckets[nr][nc] > 0:
                buckets[r][c] += 1

def new_cloud():
    for i in range(n):
        for j in range(n):
            if buckets[i][j] >= 2 and not visited[i][j]:
                buckets[i][j] -= 2
                cloud.append([i, j])

for _ in range(m):
    del_cloud = list()
    visited = [[False] * n for _ in range(n)]
    cloud = move_cloud()
    rain()
    water_copy()
    new_cloud()

print(sum(map(sum, buckets)))

설명

구현 과정은 크게 어려움이 없었고 문제에 나와있는 대로 해결이 가능했다.

하지만, 처음에 deque으로 접근했다가 시간 초과가 많이 나서 list 풀이로 변경했는데, 고치고 나서 생각해 보니 시간 초과의 문제는 다른 곳에서 났던 것 같다.

구름을 새로 생성하는 과정(new_cloud())에서 이전에 구름이 있었던 곳을 체크하는 과정에서

...
if buckets[i][j] >= 2 and [i, j] not in del_cloud:
...

이런 방식으로 확인을 했었는데 아무래도 이 부분이 시간 초과의 원인이 되었던 것 같다.
해결 코드의 방식처럼 방문 여부를 확인해서 구름을 생성하는 방식으로 바꿨더니 해결이 가능했다.

문제 해결을 위해 다른 분들의 시간 초과 원인을 살펴보니 구름을 이동하는 과정에서 좌표를 체크하는 방식의 문제가 있었는데 처음에 짰던 코드와 유사한 곳에서도 시간 초과가 날 수 있다는 사실을 알았다.

...
        while new_i0 < 0:
            new_i0 += n

        if new_i0 >= n:
            new_i0 %= n

        while new_i1 < 0:
            new_i1 += n

        if new_i1 >= n:
            new_i1 %= n

이러한 느낌의 코드도 시간 초과가 날 수 있다고 하니 유의해야겠다.