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

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

16234번: 인구 이동

---

아이디어

1. 파이썬 이슈

• 처음에 dfs로 접근했는데 Python3로 채점하면 80% 에서 시간 초과가 발생했다. (pypy로는 통과)
• 포기하고 bfs로 접근했는데 또 시간 초과가 발생해서 조건을 몇 개 추가하니 통과됨.
• c++ 였으면 dfs로 통과했을 거 같다.

 

2. 구현 파트는 크게 두 개로 나눴다.

• 국경선 열기
• 인구수 분배

---

코드 설명

기본적으로 이중 for문을 사용해서 각 위치마다 bfs를 실행한다.

여기서 방문한 적이 있거나 check함수를 통해 bfs를 실행할지 말지 결정한다.

인구이동이 있을 때까지 계속 반복해야 하므로 population_flag라는 변수를 사용해 인구변화를 표시한다.

 

def bfs(x,y):
    global population_flag
    union_list = []
    union_person = 0
    deq = deque()

    visited[x][y] = True
    union_list.append([x,y])
    union_person += country[x][y]
    deq.append([x,y])

    while deq:
        a,b = deq.popleft()
        for i in range(4):
            nx, ny = a + dx[i], b + dy[i]
            if 0 <= nx < N and 0 <= ny < N and not visited[nx][ny]:
                if L <= abs(country[a][b] - country[nx][ny]) <= R:
                    visited[nx][ny] = True
                    deq.append([nx,ny])
                    union_list.append([nx,ny])
                    union_person += country[nx][ny]

    union_len = len(union_list)

연합한 나라를 담을 리스트와 연합 인구수를 저장할 변수를 설정한다.

상, 하, 좌, 우로 인구 차이 조건을 확인하며 bfs 탐색한다.

큐에 넣을 때마다 union_list와 union_person을 업데이트해준다.

 

def check(x, y):
    for i in range(4):
        nx, ny = x + dx[i], y + dy[i]
        if 0 <= nx < N and 0 <= ny < N:
            if L <= abs(country[x][y] - country[nx][ny]) <= R:
                return 1
    return 0

이 함수는 시간 초과 때문에 추가로 만든 함수다. 연합국가가 1개라면 bfs함수를 호출할 필요가 없으므로

미리 주위 국가와 인구수 차이를 판단해 리턴을 해준다.

 

---

전체 코드

import sys
from collections import deque
N, L, R = map(int,input().split())
country = []
dx = [1,-1,0,0]
dy = [0,0,1,-1]
population_flag = True
result = 0

for _ in range(N):
    country.append([int(x) for x in sys.stdin.readline().rstrip().split()])

def bfs(x,y):
    global population_flag
    union_list = []
    union_person = 0
    deq = deque()

    visited[x][y] = True
    union_list.append([x,y])
    union_person += country[x][y]
    deq.append([x,y])

    while deq:
        a,b = deq.popleft()
        for i in range(4):
            nx, ny = a + dx[i], b + dy[i]
            if 0 <= nx < N and 0 <= ny < N and not visited[nx][ny]:
                if L <= abs(country[a][b] - country[nx][ny]) <= R:
                    visited[nx][ny] = True
                    deq.append([nx,ny])
                    union_list.append([nx,ny])
                    union_person += country[nx][ny]

    union_len = len(union_list)

    if union_len >= 2:
        for u in union_list:
            country[u[0]][u[1]] = union_person // union_len
        population_flag = True

def check(x, y):
    for i in range(4):
        nx, ny = x + dx[i], y + dy[i]
        if 0 <= nx < N and 0 <= ny < N:
            if L <= abs(country[x][y] - country[nx][ny]) <= R:
                return 1
    return 0

while population_flag:
    population_flag = False
    visited = [[False] * N for _ in range(N)]
    for i in range(N):
        for j in range(N):
            if not visited[i][j] and check(i,j):
                bfs(i,j)

    if population_flag:
        result += 1

print(result)

 

아래는 dfs로 작성한 코드

import sys
sys.setrecursionlimit(10**4)
N, L, R = map(int,input().split())
country = []
dx = [1,-1,0,0]
dy = [0,0,1,-1]

for _ in range(N):
    country.append([int(x) for x in sys.stdin.readline().rstrip().split()])
visited = [[False]*N for _ in range(N)]

def dfs(x,y):
    global union_person
    visited[x][y] = True
    union_list.append([x,y])
    union_person += country[x][y]

    for i in range(4):
        nx, ny = x + dx[i], y + dy[i]
        if 0 <= nx < N and 0 <= ny < N and not visited[nx][ny]:
            if L <= abs(country[x][y] - country[nx][ny]) <= R:
                dfs(nx,ny)

population_flag = True
result = 0

while population_flag:
    population_flag = False
    visited = [[False] * N for _ in range(N)]
    for i in range(N):
        for j in range(N):
            if not visited[i][j]:
                union_list = []
                union_person = 0
                dfs(i,j)
                union_len = len(union_list)

                if union_len >= 2:
                    population_flag = True

                for u in union_list:
                    country[u[0]][u[1]] = union_person//union_len

    if population_flag:
        result += 1

print(result)