[BFS] 경쟁적 전염 (python)

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

18405번: 경쟁적 전염

 

POINT

1. 주어진 시간 s까지 반복
2. 모든 바이러스의 순서 파악 -> 힙큐 사용
3. 번호가 낮은 종류의 바이러스부터 퍼짐 -> BFS

 

 

우선 처음 풀이는 시간 초과가 나왔다.

매 시각 바이러스의 순서를 파악하는 함수 find_virus()를 따로 만들었었는데,

바이러스를 뿌리는 함수 spread()와 만나

spread(find_virus) 하면 3차 반복문이 되어서 시간초과가 나온 것 같다.

 

=> 사실상 바이러스들을 담는 힙큐를 구할때 find_virus()를 따로 만들어주지 않고

퍼뜨릴 때마다 힙큐에 담아주어도 된다!

 

import heapq
n, k = map(int, input().split())
graph = []
for i in range(n):
    graph.append(list(map(int, input().split())))


s, x, y = map(int, input().split())

dx = [0, 0, 1, -1]
dy = [1, -1, 0, 0]


def find_virus():
    q = []
    for i in range(n):
        for j in range(n):
            if graph[i][j] != 0:
                heapq.heappush(q, (graph[i][j], (i, j)))
    return q


def spread(virus):
    global graph
    while virus:
        num, (now_x, now_y) = heapq.heappop(virus)
        for i in range(4):
            nx = now_x + dx[i]
            ny = now_y + dy[i]
            if 0 <= nx < n and 0 <= ny < n:
                if graph[nx][ny] == 0:
                    graph[nx][ny] = num


for i in range(s):
    virus = find_virus()
    spread(virus) # 사실상 spread(find_virus)이고 3차 반복문

print(graph[x-1][y-1])

 

답 풀이

바이러스를 갖는 힙큐는 따로 함수로 빼줄 필요없이, spread()에서 퍼뜨릴때마다 힙큐에 넣어주었다.

(최초 바이러스들은 graph 초기화할때 담아주었다)

 

그리고 이번 타임에서 spread할 바이러스들은 virus로 담고, (이번 타임에서 4방향으로 퍼뜨린 애들은 다음에 굳이 또 넣을 필요X)

이번에 퍼진 바이러스들은 새로운 힙큐 new_virus에 담았다.

=> 이 사이클 반복! (그림 참고)

import heapq
n, k = map(int, input().split())
graph = []
virus = []
for i in range(n):
    graph.append(list(map(int, input().split())))
    for j in range(n):
        if graph[i][j] != 0:
            heapq.heappush(virus, (graph[i][j], (i, j)))

s, x, y = map(int, input().split())

dx = [0, 0, 1, -1]
dy = [1, -1, 0, 0]



def spread(virus):
    global graph
    new_virus = []
    while virus:
        num, (now_x, now_y) = heapq.heappop(virus)
        for i in range(4):
            nx = now_x + dx[i]
            ny = now_y + dy[i]
            if 0 <= nx < n and 0 <= ny < n:
                if graph[nx][ny] == 0:
                    graph[nx][ny] = num
                    heapq.heappush(new_virus, (num, (nx, ny)))
    return new_virus


for i in range(s):
    virus = spread(virus) # 매시각 virus 담는 힙큐 초기화해주기!

print(graph[x-1][y-1])