[python] 백준 24513 : 좀비 바이러스

문제 링크
https://www.acmicpc.net/problem/24513

접근

N,M이 둘다 1000이다. 가로 세로가 1000인 경우, 맵의 총 크기는 100만으로 꽤나 크다. 딱 O(NM) 정도로 풀 수 있으면 좋을 것이다.

풀이

1번,2번 바이러스가 퍼지는 것은 그냥 일반적인 bfs 방식이다.
다음은 준비 과정이다.

input = lambda: sys.stdin.readline().rstrip()
in_range = lambda y, x: 0 <= y < n and 0 <= x < m
dy = [1, 0, -1, 0]
dx = [0, -1, 0, 1]

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

v = [0, deque(), deque()] # v[1]과 v[2]는 각각 1번, 2번 바이러스가 방문할 마을들의 좌표들을 담아놓는 큐
cnt1, cnt2, cnt3 = 1, 1, 0 # 시작 위치에 이미 바이러스 있음

for i in range(n):
    for j in range(m):
        if board[i][j] >= 1:
            v[board[i][j]].append((i, j)) # 해당 바이러스에 맞는 큐에 바이러스 시작점 집어넣음

한 턴에 새로 방문하게 될 마을들의 좌표를 tmp 라는 set에 넣으면서 탐색하는 방식으로 구현하였다. 한 턴이 끝나면 tmp는 다시 큐에 들어가 탐색을 반복한다.

while v[1] or v[2]:
    tmp1, tmp2 = set(), set()

    while v[1]:
        y, x = v[1].popleft()
        for k in range(4):
            ny, nx = y + dy[k], x + dx[k]
            if in_range(ny, nx) and board[ny][nx] == 0:
                tmp1.add((ny, nx))
    while v[2]:
        y, x = v[2].popleft()
        for k in range(4):
            ny, nx = y + dy[k], x + dx[k]
            if in_range(ny, nx) and board[ny][nx] == 0:
                tmp2.add((ny, nx))

set을 이용한 이유는 3번 바이러스 때문이다.
이번 턴에 감염시킬 마을들의 좌표를 알아야 1,2번 바이러스가 동시에 감염시키는 마을의 좌표를 알 수 있기 때문이다.

    tmp3 = tmp1 & tmp2
    for y, x in tmp3:
        board[y][x] = 3
        cnt3 += 1

    for y, x in tmp1 - tmp3:
        board[y][x] = 1
        v[1].append((y, x))
        cnt1 += 1

    for y, x in tmp2 - tmp3:
        board[y][x] = 2
        v[2].append((y, x))
        cnt2 += 1

tmp1과 tmp2에 둘 다 있는 좌표는 1,2번 바이러스가 동시에 도달한 마을이라는 뜻이고, 이는 곧 3번 바이러스가 생성되었다는 것이다.
따라서 tmp1,tmp2의 교집합인 tmp3를 만들고 보드에 3이라고 표기한다.

그 다음은 각 1,2번 바이러스에 대해서 보드에 표기하고, 새로 방문했으므로 알맞는 큐에 좌표를 집어넣어 탐색을 계속 할 수 있게 한다.

정답 코드

import sys
from collections import deque

sys.setrecursionlimit(10 ** 8)  # pypy 제출시 삭제!
input = lambda: sys.stdin.readline().rstrip()
in_range = lambda y, x: 0 <= y < n and 0 <= x < m
dy = [1, 0, -1, 0]
dx = [0, -1, 0, 1]

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

v = [0, deque(), deque()]  # v[1]과 v[2]는 각각 1번, 2번 바이러스가 방문할 마을들의 좌표들을 담아놓는 큐
cnt1, cnt2, cnt3 = 1, 1, 0  # 시작 위치에 이미 바이러스 있음

for i in range(n):
    for j in range(m):
        if board[i][j] >= 1:
            v[board[i][j]].append((i, j)) # 해당 바이러스에 맞는 큐에 바이러스 시작점 집어넣음


while v[1] or v[2]:
    tmp1, tmp2 = set(), set()

    while v[1]:
        y, x = v[1].popleft()
        for k in range(4):
            ny, nx = y + dy[k], x + dx[k]
            if in_range(ny, nx) and board[ny][nx] == 0:
                tmp1.add((ny, nx))
    while v[2]:
        y, x = v[2].popleft()
        for k in range(4):
            ny, nx = y + dy[k], x + dx[k]
            if in_range(ny, nx) and board[ny][nx] == 0:
                tmp2.add((ny, nx))

    tmp3 = tmp1 & tmp2
    for y, x in tmp3:
        board[y][x] = 3
        cnt3 += 1

    for y, x in tmp1 - tmp3:
        board[y][x] = 1
        v[1].append((y, x))
        cnt1 += 1

    for y, x in tmp2 - tmp3:
        board[y][x] = 2
        v[2].append((y, x))
        cnt2 += 1

print(cnt1, cnt2, cnt3)