[백준/파이썬]16234 인구 이동

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

16234번: 인구 이동

 

문제

백준 16234번 인구 이동

한동안 알고리즘 문제를 못 풀었었다.. 글씨도 눈에 안 들어오는 그런 기간이었다. 그런데 오랜만에 재밌는 문제를 풀었다.

 

문제를 풀기 위해서는 BFS알고리즘을 약간 응용하면 된다. 질문게시판 보니 DFS로 푸는 사람도 많은 것 같다. 

 

문제를 계속 풀다 보니 DFS보단 BFS가 더 실용적이고 더 편한 것 같다. 물론 문제에 따라 다르겠지만.

 

 

문제 해결 방법을 설명하기 위해 백준에 있는 예제 입력 1의 설명 이미지를 가져왔다.

 

다음과 같이 입력 되었을 때, 연합을 찾는 과정을 BFS 혹은 DFS로 찾아 주면 된다. 예제 입력 1의 연합이 되는 기준은 두 나라의 인구 차이가 20 이상, 50 이하 일 때 연합이 된다.

 

다음은 연합이 된 모습이다.

 

연합을 찾는 코드는 다음과 같다.

from collections import deque

def find_alliance_bfs(start_node) :
    queue = deque()
    queue.append(start_node)

    while queue :
        now_c, now_r = queue.popleft()

        for k in range(4) :
            next_c, next_r = now_c + dx[k], now_r + dy[k]

            if next_c < 0 or next_r < 0 or next_c >= n or next_r >= n :
                continue
            if visit[next_c][next_r] == 0 :
                if abs(map_[now_c][now_r] - map_[next_c][next_r]) >= l and abs(map_[now_c][now_r] - map_[next_c][next_r]) <= r :
                    alliance[alliance_idx].append([next_c, next_r])
                    visit[next_c][next_r] = 1
                    queue.append([next_c, next_r])
n, l, r = map(int, input().split())

map_ = []

for i in range(n) :
    map_.append(list(map(int, input().split())))

dx, dy = [-1, 1, 0, 0], [0, 0, -1, 1]

visit = [[0 for _ in range(n)] for _ in range(n)]
alliance = [[]for _ in range(n*n)]
alliance_idx = 0
for i in range(n) :
    for j in range(n) :
        if visit[i][j] == 0 :
            visit[i][j] = 1
            alliance[alliance_idx].append([i, j])
            find_alliance_bfs([i, j])
            alliance_idx += 1
print(alliance)

각 인덱스(ex.(0,0),(0,1),(1,0),(1,1))를 방문하지 않았다면 각 인덱스를 시작으로 BFS탐색으로 연합을 찾기 시작한다.

 

위 코드를 실행해 보면

 

[[[0, 0], [1, 0], [0, 1], [1, 1]], [], [], []]

 

이런 출력이 나오게 되는데, 입력 예제 설명 이미지와 같이 리스트 형식으로 연합이 결정된 모습을 볼 수 있다.

 

연합이 결정되는 과정을 반복해 출력을 하면

 

[[[0, 0], [1, 0], [0, 1], [1, 1]], [], [], []]
[[[0, 0]], [[0, 1]], [[1, 0]], [[1, 1]]]

 

이런 결과가 나온다.

 

2개 이상의 국가로 연합이 만들어 질 수 없다 == 더 이상 인구 이동할 수 없다

라는 뜻으로 해석할 수 있다.

 

그래서 첫 번째 예시 입력의 출력 결과는 1번의 인구이동만 있으므로 답은 1이 된다.

 

이런 식으로 문제를 풀면 된다

 

정답

from collections import deque
import sys

input = sys.stdin.readline

def find_alliance_bfs(start_node) :
    queue = deque()
    queue.append(start_node)

    while queue :
        now_c, now_r = queue.popleft()

        for k in range(4) :
            next_c, next_r = now_c + dx[k], now_r + dy[k]

            if next_c < 0 or next_r < 0 or next_c >= n or next_r >= n :
                continue
            if visit[next_c][next_r] == 0 :
                if abs(map_[now_c][now_r] - map_[next_c][next_r]) >= l and abs(map_[now_c][now_r] - map_[next_c][next_r]) <= r :
                    alliance[alliance_idx].append([next_c, next_r])
                    visit[next_c][next_r] = 1
                    queue.append([next_c, next_r])

def move_population() :
    for alliance_ in alliance :
        if alliance_ :
            alliance_population = 0
            for i in range(len(alliance_)) :
                alliance_population += map_[alliance_[i][0]][alliance_[i][1]]
            for i in range(len(alliance_)) :
                map_[alliance_[i][0]][alliance_[i][1]] = int(alliance_population / len(alliance_))

n, l, r = map(int, input().split())

map_ = []

for i in range(n) :
    map_.append(list(map(int, input().split())))

dx, dy = [-1, 1, 0, 0], [0, 0, -1, 1]

ans = 0
while True :
    visit = [[0 for _ in range(n)] for _ in range(n)]
    alliance = [[]for _ in range(n*n)]
    alliance_idx = 0
    for i in range(n) :
        for j in range(n) :
            if visit[i][j] == 0 :
                visit[i][j] = 1
                alliance[alliance_idx].append([i, j])
                find_alliance_bfs([i, j])
                alliance_idx += 1

    print(alliance)
    move_population()

    if alliance_idx == n*n :
        break
    ans += 1

print(ans)

파이썬3으로 제출하면 시간초과가 나고 pypy로 올려야 정답이 나온다 ㅜ 파이썬 시간초과까지 해결할 방법은 찾지 못했다.