백준 16234번 : 인구 이동

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

16234번: 인구 이동

유독 풀면서 실수가 많았던 문제. 풀이 자체는 다른 BFS 그래프 탐색 문제와 다르지 않다.

 크기가 NxN인 땅 안에 1x1 크기의 나라가 하나씩 존재한다. 각 나라별로 인구수가 정해져있고 인접 나라와의 인구수의 차이에 따라 인접한 국경선을 열 수 있다. 그렇게 해서 열린 국경선을 따라 하루동안 인구가 이동하는데 특정 나라에서 출발하여 방문이 가능한 나라들은 모두 연합이며, 한 연합에 속한 모든 국가들은 하루가 지나면 인구 이동으로 인해 인구가 일정한 값으로 조정된다. 인구가 더이상 이동하지 않을때까지 며칠을 이동했는지 구하는 문제이다. 

 

 아이디어는 단순하다. 존재하는 모든 나라에서 시작해서 갈 수 있는 나라들을 그래프 탐색으로 방문하면 해당 나라들은 하나의 연합이 된다. 나라를 방문할때마다 방문한 나라와 인접 나라와의 인구수를 비교하고 조건에 맞으면 방문한다. 더이상 갈 수 없을때까지 방문한 뒤, 방문했던 모든 나라들의 인구 수 총합과 좌표를 리턴하도록 하여 해당 좌표에 있는 나라들의 인구 수를 전부 조정해준다. 이미 방문한 나라는 이미 특정 연합에 속해있으므로 방문할 필요가 없다.

 

 이 과정을 더 이상 인구 이동이 일어나지 않을때까지 반복하면 되는데, 어떻게 해야 멈출 수 있을까?

위 아이디어에 따르면 모든 나라를 빠짐없이 방문해서 특정 연합에 속하도록 하는데, 연합인 나라들 끼리는 인구 이동이 반드시 일어난다. 그러니까 어떤 연합에 속해있는 나라가 둘 이상이면 그 나라 사이에서는 인구이동이 일어나고, 그 말은 즉 연합에 하나의 나라만 속해있을 때 인구이동이 멈춘다는 뜻이다. 다시 말하면 연합의 개수와 총 나라의 개수가 같으면 인구 이동이 종료된다.

 

 말했던 대로 단순하지만, 구현하면서 몇가지 실수가 있었다. 

bfs함수에서 큐를 초기화할때 초기값으로 탐색의 출발점이 되는 나라의 좌표값(튜플 자료형)을 갖도록 했는데, 쌍(a, b)으로 돼있는 좌표를 리스트의 원소가 아닌 원소가 2개인 튜플(queue = deque((a,b))로 초기화하는 바람에 오류를 찾느라 시간이 걸렸었다. 비슷한 문제를 풀다보니 손에 익어버려서 무의식적으로 작성한 것 같다.

 

추가적으로 내가 작성한 코드는 pypy 환경에서만 정답 판정을 받았다. 이론적으론 n의 최댓값이 50이므로 O(n^4) = 6,250,000 까지도 가능하지만 반복과 함수호출이 많다보니 복잡한 코드에 약한 python3 환경에서는 시간초과 판정을 받는 것 같다. 이전에 풀었던 그래프 탐색 문제 중에도, 문제를 해결할 수 있는 최소 시간 복잡도 내에서 작동하는 코드를 작성했었는데 시간초과를 판정받은 문제가 있었다. 이보다 반복을 적게하여 구현하는 방법이 있다면 참고하고싶다.

import sys
from collections import deque
input = sys.stdin.readline

def bfs(g,visit,a,b):
    dx = [1,0,-1,0]
    dy = [0,1,0,-1]
    p, u = g[a][b], [(a,b)]
    queue = deque([(a,b)])
    while queue:
        x, y = queue.popleft()
        for i in range(4):
            px, py = x+dx[i], y+dy[i]
            if px < 0 or px > n-1 or py < 0 or py > n-1:
                continue
            if l <= abs(g[x][y] - g[px][py]) <= r and not visit[px][py]:
                queue.append((px, py))
                u.append((px, py))
                visit[px][py] = True
                p += g[px][py]
    return p, u

n, l, r = map(int, input().split())
nation = [list(map(int, input().split())) for _ in range(n)]
union_cnt, cnt, people, union = 0, 0, 0, []

while union_cnt != n*n:
    visited = [[False]*n for _ in range(n)]
    union_cnt = 0
    for i in range(n):
        for j in range(n):
            if not visited[i][j]:
                visited[i][j] = True
                people, union = bfs(nation,visited,i,j)
                union_cnt += 1
                for x, y in union:
                    nation[x][y] = int(people/len(union))
    cnt += 1
print(cnt-1)