[Python] 백준 / gold / 16234번 (인구 이동)

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

BFS 알고리즘 문제이다.
unit 리스트를 생성하고, BFS 탐색을 할때마다 연합에 가입할 나라의 좌표를 담는다. 연합에 새롭게 가입할 나라가 없을 때 인구 이동이 종료된다.

visit 리스트로 방문여부를 처리하고, abs(graph[nx][ny] - graph[x][y]) 절대값 함수를 이용해서 다음 탐색할 나라와 현재 나라의 인구 수 차이를 계산한다.

Python3 으로 제출했을때는 시간초과가 났지만, Pypy3 으로 제출했을때는 통과했다..

정답 코드

import sys
from collections import deque

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

direction = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # 상하좌우

def bfs(a, b, graph, visit):
    unit = []  # 연합에 가입한 나라 좌표
    people = 0  # 나라 인구수의 합

    q = deque()
    q.append((a, b))
    unit.append((a, b))
    people += graph[a][b]
    visit[a][b] = True

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

        for d in direction:
            nx = x + d[0]
            ny = y + d[1]

            if 0 <= nx < N and 0 <= ny < N and not visit[nx][ny]:
                if L <= abs(graph[nx][ny] - graph[x][y]) <= R:
                    q.append((nx, ny))
                    unit.append((nx, ny))
                    people += graph[nx][ny]
                    visit[nx][ny] = True

    newPeople = people // len(unit)

    for a, b in unit:
        graph[a][b] = newPeople

    return True if len(unit) >= 2 else False


day = 0
while True:
    visit = [[False] * N for _ in range(N)]
    stop = True  # stop이 True면 while문 종료

    for i in range(N):
        for j in range(N):
            if not visit[i][j]:
                check = bfs(i, j, graph, visit)
                if check:
                    stop = False

    if stop:
        break
    else:
        day += 1

print(day)