[Python] 백준(18405번) - 경쟁적 전염

안녕하세요 :)

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

풀이

1. 초기 배열을 돌면서 바이러스의 위치를 찾아 각 바이러스의 위치를 나타내는 locates 배열에 넣습니다.
2. locates를 돌면서 상하좌우에 바이러스가 안퍼진 곳이라면 바이러스를 퍼뜨립니다.
3. 바이러스를 퍼뜨릴 수 있다면 new 배열에 그 위치를 추가해주고, locates[i]를 new 배열로 바꿔줍니다.
   (이미 퍼진 위치는 상하좌우 확인안해도 되니 새로퍼진곳만 new에 넣어준 다음 그 new배열을 locates[i]로 바꿔줍니다.)
4. 2, 3번 과정을 s초 동안 반복합니다.

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


def solution():
    n, k = map(int, input().split())
    arr = [list(map(int, input().split())) for _ in range(n)]
    s, x, y = map(int, input().split())
    locates = [[] for _ in range(k+1)]  # 바이러스의 위치를 넣을 배열
    for i in range(n):
        for j in range(n):
            virus_kinds = arr[i][j]
            if virus_kinds != 0:
                locates[virus_kinds].append((i, j))

    for _ in range(s):
        for i in range(1, k+1):
            new = []
            for locate in locates[i]:
                r, c = locate
                for d in range(4):
                    nr, nc = r + dx[d], c + dy[d]
                    if not (0 <= nr < n and 0 <= nc < n):
                        continue
                    if arr[nr][nc] == 0:    # 바이러스가 안퍼진 곳이라면
                        arr[nr][nc] = i
                        new.append((nr, nc))
            locates[i] = new

    return arr[x-1][y-1]


if __name__ == '__main__':
    print(solution())