[BOJ/Python] 16234. 인구 이동

📌 문제

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

16234번: 인구 이동

 

📌 문제 접근 방법

1. 우선, 첫 번째 인구 이동이 일어나는 과정을 먼저 구현해보았다.
   방문한 적이 없는 곳이면 bfs(find_unit 함수)를 통해 연합을 구성하고 
   → 연합의 길이가 1보다 크다면 인구이동이 일어나므로 move함수를 통해 인구를 이동시킨다.
2. 위의 코드를 이용해서 인구 이동이 더 이상 일어날 수 없을 때까지 인구 이동을 진행하고,
   인구 이동이 일어날 때 마다 카운트를 해줘야 한다.
   → 연합의 길이가 1보다 큰 경우를 units 리스트에 넣고 리스트가 비어있다면 while문 종료
   → 인구 이동이 일어난 경우 move함수 호출하고 카운팅
   → 81% 에서 시간초과
3. 여러가지 시도를 해보다가 결국 pypy로 바꿔서 제출했더니 성공했다...!

 

📌 코드

# 백준 16234 인구 이동
from collections import deque

def move(unit, total):   
    for i, j in unit:
        country[i][j] = total // len(unit)


def find_unit(i, j):
    unit = [(i, j)]
    total = country[i][j]

    queue = deque()
    queue.append((i, j))

    dxy = [(1, 0), (-1, 0), (0, 1), (0, -1)]

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

        for dx, dy in dxy:
            nx, ny = x+dx, y+dy

            if 0 <= nx < N and 0 <= ny < N and not visited[nx][ny]:
                diff = abs(country[x][y] - country[nx][ny])
                if L <= diff <= R:
                    visited[nx][ny] = 1
                    queue.append((nx, ny))

                    unit.append((nx, ny))
                    total += country[nx][ny]
    
    return unit, total


N, L, R = map(int, input().split())
country = [list(map(int, input().split())) for _ in range(N)]

cnt = 0
while True:
    visited = [[0 for _ in range(N)] for _ in range(N)]
    units = []

    for i in range(N):
        for j in range(N):
            if not visited[i][j]:
                visited[i][j] = 1
                unit, total = find_unit(i, j)

                if len(unit) > 1:
                    units.append((unit, total))

    if units:
        for unit, total in units:
            move(unit, total)
        cnt += 1
    else:
        break

print(cnt)