[Python] 백준 16234번 - 인구 이동

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

16234번: 인구 이동

처음에는 모든 좌표를 방문하면서 오른쪽과 아래방향의 나라만 비교하면 중복없이 연합 국가들을 찾을 수 있겠다고 생각했다.

그렇게 연합국들을 찾아 union 변수에 다 넣고 처리해주면 풀 수 있을 줄 알았다.

실제로 예제 1~4는 해당 방식으로 문제가 풀렸는데, 예제 5에서 오류가 떴다.

 

문제는 연합국들이 2개 이상의 지역으로 나눠지는 경우였다.

연합국 지역이 2개 이상이면, 매번 union 변수를 새로 생성해서 각 지역별로 인구 계산을 해주어야 한다.

따라서 dfs를 사용하여 지역별 연합국을 탐색할 수 있도록 해주었다.

 

완성코드는 다음과 같다.

import copy

n, l, r = map(int, input().split())
board = [list(map(int, input().split())) for _ in range(n)]
count = 0
d = [(1, 0), (-1, 0), (0, -1), (0, 1)]

def dfs(x, y):
    union[(x, y)] = board[x][y]
    visited[x][y] = 1
    for i in range(4):
        nx = x + d[i][0]
        ny = y + d[i][1]
        if 0 <= nx < n and 0 <= ny < n and not visited[nx][ny]:
            if l <= abs(board[x][y] - board[nx][ny]) <= r:
                dfs(nx, ny) 
    return

while True:
    count += 1
    temp_board = copy.deepcopy(board)
    visited = [[0]*n for _ in range(n)]
    
    for x in range(n):
        for y in range(n):
            if not visited[x][y]:
                union = {}
                dfs(x, y)
                
                # 연합이 없으면 다음 연합 찾기
                if len(union) == 1:
                    continue
                
                # 연합이 있을때, 인구 수 구하기
                num = 0              
                for i in union:
                    num += union[i]
                num = num // len(union)

                # 인구 이동하기
                for key in union:
                    board[key[0]][key[1]] = num
    
    # 인구 변화가 없는 경우(연합이 없는 경우) 종료      
    if board == temp_board:
        break
    
print(count-1)