백준 인구 이동 코드 및 해설 (파이썬)

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

16234번: 인구 이동

 

* Python으로는 시간 초과, PyPy3로는 통과 

 

find_union 함수로는 현재 국가들 상태에서 국경 열 수 있는 국가들을 구합니다. 반환하는 것은 list of list인 unions로, 같은 연합국인 국가들의 좌표를 하나의 리스트로 저장합니다. 만약 연결되지 않는 연합국 그룹이 2개라면, unions는 2개의 리스트를 가진 리스트입니다. 

모든 좌표에 대해 아직 방문하지 않았으면 다음을 실행합니다. 

현재 좌표를 큐에 넣고, 큐가 빌 때까지 다음을 반복합니다. 큐를 pop하고, 이렇게 얻은 좌표의 상하좌우 좌표를 구합니다. 이 좌표 중 0~N 안에 있고 아직 방문하지 않은 곳은, pop한 좌표(cur)와 이 좌표(p)의 인구 수 차이를 구해 diff에 저장합니다. 이 diff가 L 이상, R 이하면 방문 처리를 하고 큐에 이(p)를 추가합니다. 이후 unions에 이 정보를 저장하는데, 만약 pop했던 좌표(cur)가 이미 unions의 하위 리스트에 존재한다면, 그 리스트에 이 좌표(p)를 추가합니다. 만약 pop했던 좌표(cur)가 unions의 어디에도 없다면 unions에 새로운 리스트를 추가하는데, 이 리스트에는 pop했던 좌표(cur)와 이 좌표(p)가 들어가 있습니다.

 

이 함수를 이용해 최종 답을 구합니다. 우선 처음 입력에 따라 find_union을 구합니다. 이후 unions가 빈 리스트가 될 때까지 계속해서 다음을 반복합니다. 인구 이동에 걸린 일 수 days를 하나씩 올려주고,  unions 안의 각 연합국 그룹에 대해 인구 수를 갱신합니다. 이후 visited를 초기화하고 다시 find_unions를 실행합니다.

while 문을 빠져 나온 후 days를 최종 답으로 출력합니다.

# python 시간 초과, pypy3 통과
import sys
from collections import deque

N, L, R = map(int, sys.stdin.readline().split())
world = []
for _ in range(N):
    world.append(list(map(int, sys.stdin.readline().split())))
visited = [[False] * N for _ in range(N)]


# 상하좌우
UDLR = [(-1, 0), (1, 0), (0, 1), (0, -1)]

# 국경선 열리는 나라들 구하기 
def find_union(world, N, L, R, unions):
    # 모든 좌표에 대해 
    for i in range(N):
        for j in range(N):
            # 아직 방문하지 않은 경우 
            if not visited[i][j]:
                q = deque([(i,j)])
                while q:
                    cur = q.popleft()
                    # 현재 좌표의 상하좌우
                    pos = []
                    for udlr in UDLR:
                        a,b = cur
                        x,y = udlr
                        pos.append((a+x, b+y))
                    for p in pos:
                        # 0~N 안에 있고 아직 방문하지 않은 경우 
                        if p[0] in range(N) and p[1] in range(N) and not visited[p[0]][p[1]]:
                            # cur와 상하좌우의 인구 수 차이 
                            diff = abs(world[cur[0]][cur[1]] - world[p[0]][p[1]])
                            if L <= diff and diff <= R:
                                # 방문 후 큐에 추가 
                                visited[p[0]][p[1]] = True
                                q.append(p)
                                already = False
                                # cur가 이미 다른 나라와 union인 경우, 그 union에 p 추가 
                                for num, union in enumerate(unions):
                                    if (i,j) in union:
                                        unions[num] += [p]
                                        already = True
                                # cur의 첫 union인 경우 cur, p를 unions 안의 새로운 list로 추가 
                                if not already:
                                    unions.append([cur, p])
    return unions

days = 0
# unions: list of list
# ex. (0,0)과 (1,0)이 국경을 열고 (2,1), (2,2)가 국경을 연다면, unions = [[(0,0),(1,0)], [(2,1),(2,2)]]
unions = find_union(world, N, L, R, unions=[])

# 국경을 열 수 있는 상태면 계속 반복 
while unions:
    # 인구 이동 날짜 더하기 
    days += 1
    # 각 연합국들에 대해 
    for union in unions:
        # 중복 제거 
        union = list(set(union))
        # 연합국의 총 인구 수 
        sum = 0
        for pos in union:
            sum += world[pos[0]][pos[1]]
        # 연합국의 새로운 인구 수 
        new_val = int(sum / len(union))
        # 연합국 인구 수 갱신
        for pos in union:
            world[pos[0]][pos[1]] = new_val

    # visited 초기화 
    visited = [[False] * N for _ in range(N)]
    unions = find_union(world, N, L, R, unions=[])

print(days)