백준 16234번 : 인구 이동 in Python

16234번: 인구 이동

 

www.acmicpc.net

이 문제는 각 나라의 인구수가 주어졌을 때 인구 이동이 몇 번 발생하는지 구하는 문제이다. 인구 이동의 조건은 위와 같다.

이 문제를 해결하기 위해서는 대략적으로 다음과 같은 과정을 거쳐야 한다. 우선 각 나라에 대해 그 나라를 포함하는 나라들의 연합을 만들고, 연합 내 모든 국가들의 인구 수를 모든 국가들의 인구 수의 평균으로 바꿔야 한다. 이를 방문하지 않은 모든 나라에 대해 수행하고 만약 인구 이동이 발생하지 않은 경우 위 과정을 끝내고, 인구 이동이 발생한 경우 다시 위 과정을 반복한다.

좀 더 구체적으로 위 과정을 구현하자. 우선 특정 나라의 위치가 주어질 때 그 나라를 포함하는 연합을 구하고 연합 내 나라들의 인구 수를 바꾸는 함수 BFS를 만든다. 이름에서도 알 수 있듯이 여기서는 너비 우선 탐색을 구현해 그래프를 탐색한다. 다른 부분은 일반적인 BFS 구현과 유사하나, 추가로 연합에 포함되는 나라를 저장할 배열인 visited_c와 연합 내 나라의 인구 수의 합과 연합 내 나라 수를 저장하는 변수인 population과 cnt를 추가한다. 그리고 어떤 나라와 인접하면서 조건을 만족하는 나라를 queue에 넣는 과정에서 visited_c에도 나라를 넣고 population에는 그 나라의 인구 수, cnt는 1 증가시킨다.

그리고 너비 우선 탐색을 끝낸 후 population을 cnt로 나눠서 연합 내 나라의 바뀐 인구 수 avg를 구하고 visited_c에 저장된 나라의 인구 수를 avg로 바꾼다. 그리고 함수 내에서 인구 이동이 발생했는지를 확인하기 위해 res라는 변수를 추가해 초기값으로 False를 준 뒤, 두 나라 이상의 연합이 발생한 경우 res 변수를 True로 바꾼다. 그리고 함수가 끝날 때 res를 반환함으로써 인구 이동이 발생했는지를 알려준다.

이제 BFS 함수를 이용해 인구 이동이 총 몇 번 발생하는지 구하기 위해, while loop을 돌면서 우선 res 변수와 N*N 크기의 visited 행렬을 각각 False와 모든 entry가 0이 되도록 초기화한다. 그리고 모든 나라를 순회하면서 방문하지 않은 나라가 있을 경우 해당 나라에 대해 BFS 함수를 실행한 뒤 반환된 결과와 res를 or 연산을 통해 계산한 결과를 다시 res에 저장해 한 번이라도 BFS의 반환 값이 True가 된 경우, 즉 한 번이라도 인구 이동이 발생한 경우 res가 True가 되도록 한다. 이제 res가 False인 경우 인구 이동이 한 번도 발생하지 않았으므로 while loop을 빠져나오고, 그렇지 않은 경우 인구 이동 횟수를 1 증가시키고 다시 while loop을 반복한다. 이를 실행하면 인구 이동이 더 이상 발생하지 않을 때까지의 총 인구 이동 횟수를 구할 수 있다.

이를 코드로 구현하면 다음과 같다.

from collections import deque
N, L, R = map(int, input().split())
ground = []
for _ in range(N):
    ground.append(list(map(int, input().split())))
dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]
# 총 인구 이동 횟수
total_cnt = 0
def BFS(i, j):
    global N, L, R
    # 인구 이동이 발생했는지를 나타내는 변수
    res = False
    visited[i][j] = 1
    # 인자로 주어진 나라를 포함하는 연합의 나라들을 저장하는 배열
    visited_c = []
    visited_c.append([i, j])
    queue = deque()
    queue.append([i, j])
    # 연합 내 나라들의 인구 수의 합을 저장하는 변수
    population = ground[i][j]
    # 연합 내 나라들의 수를 저장하는 변수
    cnt = 1
    while queue:
        x, y = queue.popleft()
        for i in range(4):
            next_x = x + dx[i]
            next_y = y + dy[i]
            # 인접한 나라 중 방문하지 않았으면서 국경선을 열기 위한 조건을 만족하는지 검사한다.
            if 0<=next_x<N and 0<=next_y<N and visited[next_x][next_y] == 0 and L<=abs(ground[x][y] - ground[next_x][next_y])<=R:
                visited[next_x][next_y] = 1
                # 연합에 해당 나라를 포함시킨다.
                visited_c.append([next_x, next_y])
                queue.append([next_x, next_y])
                population += ground[next_x][next_y]
                cnt += 1
                # 인구 이동이 발생하므로 res를 True로 바꾼다.
                res = True
    avg = population//cnt
    # 연합 내 인구 수를 평균 인구 수로 바꾼다.
    for c in visited_c:
        ground[c[0]][c[1]] = avg
    return res
while True:
    res = False
    visited = [[0]*N for _ in range(N)]
    # 모든 나라에 대해 순회하면서 방문하지 않은 나라에 대해 BFS를 수행한다.
    for i in range(N):
        for j in range(N):
            if visited[i][j] == 0:
                # 인구 이동이 한 번이라도 발생하면 res의 값이 True가 되고, 
                # 그렇지 않은 경우 res는 False가 된다.
                res = BFS(i, j) or res
    if not res: break
    else:
        total_cnt += 1
print(total_cnt)