파이썬으로 풀어보는 백준 16234번: 인구 이동 (삼성 A형 기출 문제)

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

16234번: 인구 이동

 

내 풀이:

def BFS(y, x):
    global TF
    q = [[y, x]]
    s = [[y, x]]
    tot = arr[y][x]
    while q:
        current = q.pop(0)
        y = current[0]
        x = current[1]
        for dir in range(4):
            ny = y + dy[dir]
            nx = x + dx[dir]
            if 0 <= ny <= N-1 and 0 <= nx <= N-1 and L <= abs(arr[y][x] - arr[ny][nx]) <= R and not visited[ny][nx]:
                visited[ny][nx] = True
                q.append([ny, nx])
                s.append([ny, nx])
                tot += arr[ny][nx]
 
    if len(s) == 1:
        return
 
    TF = True
    avg = tot // len(s)
 
    for i in s:
        arr[i[0]][i[1]] = avg
 
 
N, L, R = map(int,input().split())
arr = [list(map(int,input().split())) for _ in range(N)]
 
dy = [0, 0, -1, 1]
dx = [-1, 1, 0, 0]
cnt = -1
TF = True
while TF:
    cnt += 1
    visited = [[False] * N for _ in range(N)]
    TF = False
    for i in range(N):
        for j in range(N):
            if not visited[i][j]:
                visited[i][j] = True
                BFS(i,j)
 
print(cnt)

PyPy3, 1336ms

Python3로 채점을 했더니 시간 초과가 나서 PyPy3로 채점하여 정답처리를 받았다. BFS를 이용하여 국경을 열어서 인구가 이동하는 구역을 체크하고 평균으로 만든 다음 visited 처리를 하여 문제를 풀었다. BFS 문제 중에서는 어렵지 않은 편인 것 같은데 최적화를 잘해야 시간 초과가 나지 않도록 설계된 문제 같다. 결과적으로 Python3로 채점을 받지 못했기 때문에 최적화가 더 필요할 것 같다.