[Python] 백준(BOJ) 16234번: 인구 이동 - bfs

문제

16234번: 인구 이동

풀이

결국 시간초과가 나서 pypy로 제출했다.

import sys
from collections import deque

N, L, R = map(int, sys.stdin.readline().split(" "))

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

land = []
for _ in range(N):
    land.append(list(map(int, sys.stdin.readline().split(" "))))

t = 0
while True:
    visited = [[0 for _ in range(N)] for _ in range(N)]
    count = 0
    take = 0

    for i in range(N):
        for j in range(N):
            if visited[i][j] == 0:
                count += 1
                add = 0
                c = 1
                union = []
                queue = deque()

                queue.append([i, j])
                visited[i][j] = count
                add += land[i][j]
                union.append([i, j])

                while queue:
                    x, y = queue.popleft()

                    for k in range(4):
                        ndr = x + dr[k]
                        ndc = y + dc[k]

                        if 0 <= ndr < N and 0 <= ndc < N and visited[ndr][ndc] == 0:
                            # 차이가 범위 안에 들어오면
                            if L <= abs(land[x][y] - land[ndr][ndc]) <= R:
                                queue.append([ndr, ndc])
                                visited[ndr][ndc] = count
                                add += land[ndr][ndc]
                                union.append([ndr, ndc])
                                c += 1
                                take = 1

                for each in union:
                    land[each[0]][each[1]] = add // c

    if take == 0:
        break

    t += 1

print(t)