[백준] 16234 인구 이동 python(pypy3)

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

16234번: 인구 이동

import sys
from collections import deque

def bfs(r, c):
    visit[r][c] = 1
    q = deque()
    q.append((r, c))
    ret = [(r,c)]
    while q:
        y, x = q.popleft()
        for i in range(4):
            yy = y + dy[i]
            xx = x + dx[i]
            # 주어진 나라의 크기를 넘는 경우
            if yy < 0 or yy >= N or xx < 0 or xx >= N: continue
            # 방문한 경우
            if visit[yy][xx]: continue
            # 인구 차이가 안나는 경우
            tmp = abs(N_map[yy][xx] - N_map[y][x])
            if tmp < L or tmp > R: continue
            visit[yy][xx] = 1
            q.append((yy, xx))
            ret.append((yy, xx))
    return ret


global N, L, R
N, L, R = map(int, sys.stdin.readline().split())
N_map = [list(map(int, sys.stdin.readline().split())) for _ in range(N)]
answer = 0
# r d l u
dy = [0, 1, 0, -1]
dx = [1, 0, -1, 0]

is_change = True
while is_change:
    is_change = False

    # 1. 연합(국경 개방하는 나라) 찾기
    visit = [[0 for _ in range(N)] for _ in range(N)]
    nation = []
    for i in range(N):
        for j in range(N):
            if visit[i][j]: continue
            nation.append(bfs(i, j))

    # 2. 국경 이동
    for i in range(len(nation)):
        population = 0
        for j in range(len(nation[i])):
            population += N_map[nation[i][j][0]][nation[i][j][1]]
        post_population = population // len(nation[i])
        while nation[i]:
            y, x = nation[i].pop()
            if N_map[y][x] == post_population: continue
            is_change = True
            N_map[y][x] = post_population

    if is_change:
        answer += 1

print(answer)

 

1. 인접하는 나라의 인구차이가 L이상 R이하일때 국경을 개방하고, 모든 국경선을 열었을때 인구이동을 시작한다.

2. 이때 연합(국경선이 열려있을때 통행할수 있는 나라들)의 인구수를  (연합의 인구수 / 나라의 개수)로 갱신한다.(소수점은 버린다)

3. 다 갱신한 이후에는 국경선을 닫는다.(코드에선 구현하지 않음)

인구이동이 일어나지 않을때까지 1~3을 반복하는데 이때 인구이동이 몇번 발생하는지 출력하면된다.

# 1 연합찾기

이 문제는 그래프 탐색 문제인데 연합을 찾기 위해 너비 우선 탐색을 사용했다. 범위(N*N)를 벗어나는 경우, 방문했던 나라인 경우, 그리고 인구차이가 L보다 작거나 R보다 큰 경우를 조건으로 걸고 넘어갔다.

# 2 국경이동

찾은 연합 내에서 국경을 이동후 인구를 계산하기 위해 연합의 인구수를 구했고 연합의 크기를 구해서 나눠 주었다. 그리고 해당 나라들을 모두 순회하며 인구를 갱신해주었다. 한번이라도 갱신이 되면 answer의 값을 하나 증가시켰다.

한번도 갱신되지 않은 경우엔 움직임이 없는 것으로 판단하여 반복문을 빠져나오게 된다.

 

---

bfs 안에 국경이동하는 것까지 짜면 통과된다.

 

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

def bfs():
    is_change = False
    visit = [[0 for _ in range(N)] for _ in range(N)]
    for i in range(N):
        for j in range(N):
            if not visit[i][j]:
                visit[i][j] = 1
                ret = [(i,j)]
                q = deque([(i,j)])
                population = N_map[i][j]
                while q:
                    y, x = q.popleft()
                    for dy, dx in [(1, 0), (0, 1), (-1, 0), (0, -1)]:
                        yy = y + dy
                        xx = x + dx
                        # 주어진 나라의 크기를 넘는 경우, 방문한 경우
                        if 0 <= yy < N and 0 <= xx < N and not visit[yy][xx]:
                            # 인구 차이가 안나는 경우
                            if L <= abs(N_map[yy][xx] - N_map[y][x]) <= R:
                                visit[yy][xx] = 1
                                population += N_map[yy][xx]
                                q.append((yy, xx))
                                ret.append((yy, xx))
                if len(ret) > 1:
                    is_change = True
                    population = population // len(ret)
                    for x, y in ret:
                        N_map[x][y] = population
    return is_change


N, L, R = map(int, input().split())
N_map = [list(map(int, input().split())) for _ in range(N)]
answer = 0
while True:
    if bfs():
        answer += 1
    else:
        break

print(answer)