16234번 인구 이동(python3)

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

16234번: 인구 이동

dfs, bfs알고리즘 문제다.

내 풀이코드는 다음과 같다.

찾아보니 다른 사람들은 bfs를 많이 적용했는데, 나는 dfs를 연습하는 김에 dfs를 사용했다.

import sys

sys.setrecursionlimit(10**6)

N, L, R = map(int, input().split())
data = []
for i in range(N):
    data.append(list(map(int, input().split())))
data1 = [[0]*N for i in range(N)]
flag = True
union = []
def dfs(data, data1, x, y):
    if data1[x][y] == 1:
        return
    union.append([x, y])
    data1[x][y] = 1
    if x-1 >= 0 and L <= abs(data[x][y] - data[x-1][y]) <= R and data1[x-1][y] == 0:
        dfs(data, data1, x-1, y)
    if x+1 < N and L <= abs(data[x][y] - data[x+1][y]) <= R and data1[x+1][y] == 0:
        dfs(data, data1, x+1, y)
    if y-1 >= 0 and L <= abs(data[x][y] - data[x][y-1]) <= R and data1[x][y-1] == 0:
        dfs(data, data1, x, y-1)
    if y+1 < N and L <= abs(data[x][y] - data[x][y+1]) <= R and data1[x][y+1] == 0:
        dfs(data, data1, x, y+1)

dx, dy = [-1, 1, 0, 0], [0, 0, -1, 1]
def dfs2(data, data1, x, y):
    if data1[x][y] == 1:
        return
    union.append([x, y])
    data1[x][y] = 1
    for i in range(4):
        nx, ny = x + dx[i], y + dy[i]
        if 0 <= nx < N and 0 <= ny < N:
            if L <= abs(data[x][y] - data[nx][ny]) <= R:
                dfs2(data, data1, nx, ny)


ans = 0
while True:
    flag = False
    data1 = [[0] * N for i in range(N)]
    dx, dy = [-1, 1, 0, 0], [0, 0, -1, 1]
    for i in range(N):
        for j in range(N):
            if data1[i][j] != 1:
                dfs(data, data1, i, j)
                if len(union) > 1:
                    flag = True
                    sum = 0
                    for k in union:
                        sum += data[k[0]][k[1]]
                    sum = sum // len(union)
                    for k in union:
                        data[k[0]][k[1]] = sum
            union = []
    if not flag:
        break
    ans += 1
print(ans)

dfs를 인구수의 차이가 L, R 사이라는 조건을 만족하며 dfs를 적용해주고 인구를 분배한다.

이를 더이상 인구분배가 일어나지 않을 때까지 반복한다. 백준에 문제를 제출하는데 python3로 제출하면 재귀함수 깊이 오류가 발생했고 pypy3로 제출하면 정답으로 처리됐다. 이유를 찾아보니 python3는 기본적으로 재귀함수의 깊이가 최대 1000으로 설정돼있기 때문이었다.

import sys

sys.setrecursionlimit(10**6)

코드를 추가해서 최대 재귀함수의 깊이를 크게 설정하니 python3로도 정답처리가 됐다.