백준 18405번: 경쟁적 전염

풀긴 풀었는데.. 동작 시간이 영 구리게 나와서 다른 실력자분들의 코드를 보며 공부했고, 덕분에 많이 배울 수 있었던 고마운 문제다. 각설하고 처음 작성해서 통과하긴 했던 내 코드

from collections import deque
import heapq
import sys
input = sys.stdin.readline

N, K = list(map(int, input().split()))
flask = [list(map(int, input().split()))for _ in range(N)]
S, X, Y = list(map(int, input().split()))

def effectVirus(virusNum, x, y):
    dx = [-1,1,0,0]
    dy = [0,0,-1,1]
    for i in range(4):
        if N-1 < x+dx[i] or x+dx[i] < 0 or N-1 < y+dy[i] or y+dy[i] < 0:
            continue
        if flask[x+dx[i]][y+dy[i]] == 0:
            flask[x+dx[i]][y+dy[i]] = virusNum

def checkVirus():
    result = []
    for i in range(N):
        for j in range(N):
            if flask[i][j] != 0:
                result.append((flask[i][j], i, j))
    return result

second = 0
while second < S:
    viruses = checkVirus()
    if len(viruses) == N**2:
        break
    heapq.heapify(viruses)
    while viruses:
        virusnum, x, y = heapq.heappop(viruses)
        effectVirus(virusnum, x, y)
    second += 1

print(flask[X-1][Y-1])

특이점은 무식하게 풀었다는 것? 😂 우선순위가 높은 바이러스부터 퍼뜨려야 하기 때문에, 힙을 써봐야 겠다는 생각만 하고 그냥 무작정 코드를 짰던 것 같다.

1. 매 초 시험관 전체를 탐색하며 바이러스를 찾아내고 (불필요😅)
2. 바이러스라고 판단된 모든 녀석들(불필요😅)의 상하좌우를 감염시킨다.
3. 감염을 완전히 마치거나, 원하는 시간이 된다면 반복을 멈춘다.

지금 다시 보면 이게 어떻게 통과됐나 싶다.. 부끄럽다 부끄러워
실력자들의 코드를 보고 다시 작성한 코드는 다음과 같다.

from collections import deque
import heapq
import sys
input = sys.stdin.readline

N, K = list(map(int, input().split()))
flask = [list(map(int, input().split()))for _ in range(N)]
S, X, Y = list(map(int, input().split()))
dx = [-1,1,0,0]
dy = [0,0,-1,1]

virusQ = []
for i in range(N):
    for j in range(N):
        if flask[i][j]:
            heapq.heappush(virusQ, (0,flask[i][j],i,j))
            
while virusQ:
    sec, virusNum, x, y = heapq.heappop(virusQ)
    for i in range(4):
        xx = x + dx[i]
        yy = y + dy[i]
        if 0<=xx<=N-1 and 0<=yy<=N-1 and flask[xx][yy] == 0 and sec<S:
            flask[xx][yy] = virusNum
            heapq.heappush(virusQ, (sec+1,virusNum,xx,yy))

print(flask[X-1][Y-1])

가장 크게 배운 점은, 튜플들을 담고 있는 리스트를 heap화 할때, 첫 인자 기준으로만 heapify해주는게 아니라 두번째 이후까지도 고려한 heap화를 해준다는 것. 아니 이게 쓰면서도 뭔 소린지 헷갈리는데 예를 들어 (시간, 바이러스 넘버)와 같은 튜플을 heapify해서 처리하면, 시간이 작은 녀석을을 최우선으로 처리하고, 이후에는 바이러스 넘버가 작은 녀석들을 차순으로 처리할 수 있다는 뜻이다. 이를 이용한 게 위의 코드인데, sec을 튜플의 가장 왼쪽, 그 바로 뒤에 virusNum을 두어 내가 원하는 순서대로 깔끔하게 태스크를 처리할 수 있었다. 여기에 시간을 큐에 추가로 넣어주어 통째로 BFS로 처리할 수 있는건 덤. 이러한 개선을 통해 통과 시간을 768ms -> 188ms로 크게 개선할 수 있었다.