[그래프 순회] 백준 16234: 인구 이동

문제: N*N 크기의 땅이 있다. 각각의 땅에는 나라가 하나씩 존재하고, r행 c열엔 A[r][c]명이 살고 있다.

인접한 나라 사리에는 국경선이 존재한다. 인구 이동은 다음과 같은 방식으로 진행된다.

 

1. 국경선을 공유하는 두 나라의 인구 차이가 L명 이상 R명 이하라면 두 나라가 공유하는 국경선을 하루동안 연다.

2. 위의 조건에 의해 열어야하는 국경선이 모두 열렸다면, 인구 이동을 시작한다.

3. 국경선이 열려있어 인접한 칸만을 이용해 이동할 수 있으면, 그 나라를 오늘 하루 동안 연합이라 한다.

4. 연합을 이루고 있는 각 칸의 인구수는 연합의 인구 수 / 연합을 이루고 있는 칸의 수가 된다. 소수점 버림.

5. 연합을 해체하고 모든 국경선을 닫는다.

 

위의 과정을 인구 이동이 없을 때까지 반복한다.

각 나라의 인구수가 주어졌을 때, 인구 이동이 며칠 동안 발생하는지 구하는 프로그램을 작성하라.

---

입력: 첫째 줄에 N, L, R 주어짐.

둘째 줄부터 N개의 줄에 각 나라의 인구수가 주어짐. r행 c열에 주어지는 정수는 A[r][c]의 값.

인구 이동이 발생하는 일수가 2000번 보다 작거나 같은 입력만 주어짐.

---

 

얘도 단지 번호 메기기와 비슷한 문제 같다. 로직을 간단히 세워보았다.

1. 인구의 분포를 arr라는 이차원 리스트에 저장하고, arr[0][0] 에서 순회를 시작하자.

2. 델타탐색을 이용, 국경을 열 수 있는 곳은 방문처리와 병합처리를 해주자. (visited 리스트와 beonghaped 리스트를 활용하였다.)

3. arr[0][0] 지점에서 시작해 인구 이동이 가능한 모든 지점을 순회했다면, beonghaped 리스트에 값이 1로 들어가 있는 좌표와 같은 좌표에 해당하는 arr리스트의 값을 모두 더한 후 그 좌표값을 beonghaped 리스트에 값이 1로 들어가 있는 좌표의 수로 나누어준다. 나머지는 버림처리.

4. 3번까지 수행하였으면 beonghaped 리스트 내의 값을 모두 0으로 바꾸어주고, visited 리스트의 값이 0이 아닌 다른 좌표에서 순회를 다시 시작하자.

5. arr리스트에 대한 탐색이 1회 완료되었을 경우 ans 값을 1 증가시킨다.

6. arr리스트에 대한 탐색을 다시 시작하고, 해당 탐색이 진행되는 동안 flag값(병합의 유무를 판단함)의 변동에 따라 탐색을 계속할지 말지 결정한다.

 

위의 로직을 활용해 코드를 짜 보았다.

 

from collections import deque

# 델타탐색
di = [0, 0, 1, -1]
dj = [1, -1, 0, 0]

def bfs(k, l):
    global flag
    count = 1
    plus = arr[k][l]
    beonghaped[k][l] = 1
    visited[k][l] = 1
    while queue:
        start = queue.popleft()
        for i in range(4):
            ni = start[0] + di[i]
            nj = start[1] + dj[i]
            if 0 <= ni < N and 0 <= nj < N and visited[ni][nj] == 0: # 범위 내에 있으면서 아직 방문하지 않았다면
                if L <= ((arr[start[0]][start[1]] - arr[ni][nj]) ** 2) ** (1/2) <= R and beonghaped[ni][nj] == 0:
                    count += 1
                    plus += arr[ni][nj]
                    beonghaped[ni][nj] = 1
                    visited[ni][nj] = 1
                    queue.append([ni, nj, arr[ni][nj]])
                    flag = 1 # 병합 요소가 추가로 있으면 flag값을 1로 바꾸자

    return count, plus

N, L, R = map(int, input().split())
arr = []
for i in range(N):
    arr.append(list(map(int, input().split())))

# 병합 가능한 국가의 모든 국경선을 열어두어야 한다.
queue = deque()
visited = [[0] * N for _ in range(N)]
beonghaped = [[0] * N for _ in range(N)] # 병합 유무 판단

# 덩어리를 바꿔주는 역할
flag = 1
ans = 0
while flag:
    flag = 0
    ans += 1
    for i in range(N):
        for j in range(N):
            if visited[i][j] == 0: # 아직 탐색하지 않은 덩어리라면
                queue.append([i, j, arr[i][j]])
                count, plus = bfs(i, j)
                for k in range(N):
                    for z in range(N):
                        if beonghaped[k][z] == 1: # 병합 대상이면
                            arr[k][z] = plus // count

                beonghaped = [[0] * N for _ in range(N)] # 한 덩어리 병합 끝났으면 병합 정보 초기화 해야 함
    visited = [[0] * N for _ in range(N)]

print(ans - 1) # 최초 실행시 병합 요소가 없더라도 한번은 더해지기 때문에, 1을 빼주어야 한다.

제출해봤는데 80퍼센트에서 시간초과가 뜬다. 5중 반복문이 코드에 포함되어 있어 그런 것 같다. 어디를 수정해야 할까? 곰곰히 고민해 봤다. 고민해본 결과 현재의 로직으로는 5중 반복문을 해결할 수 없을 것 같다는 결론이 나왔다.

1. 시작점을 찾기 위해 2중 for문이 필수

2. 인구 변화가 며칠 일어나는지 체크하기 위한 while문 필수

3. bfs에 포함되는 while문 1개와 for문 1개 필수

 

그럼 어디에서 시간을 줄여야 할까? 병합 예정인 좌표를 하나의 리스트에 메모해두고, 그 메모 결과를 순회하며 arr값을 plus // count로 바꾸어 주면 될 것 같다는 생각이 들었다.

 

수정한 코드는 다음과 같다.

import sys
sys.stdin = open('16234input.txt')
from collections import deque

# 델타탐색
di = [0, 0, 1, -1]
dj = [1, -1, 0, 0]

def bfs(k, l):
    global flag
    count = 1
    plus = arr[k][l]
    # beonghaped[k][l] = 1
    beonghaped.append([k, l])
    visited[k][l] = 1
    while queue:
        start = queue.popleft()
        for i in range(4):
            ni = start[0] + di[i]
            nj = start[1] + dj[i]
            if 0 <= ni < N and 0 <= nj < N and visited[ni][nj] == 0: # 범위 내에 있으면서 아직 방문하지 않았다면
                if L <= ((arr[start[0]][start[1]] - arr[ni][nj]) ** 2) ** (1/2) <= R and [ni, nj] not in beonghaped:
                    count += 1
                    plus += arr[ni][nj]
                    beonghaped.append([ni, nj])
                    visited[ni][nj] = 1
                    queue.append([ni, nj, arr[ni][nj]])
                    flag = 1 # 병합 요소가 추가로 있으면 flag값을 1로 바꾸자

    return count, plus

N, L, R = map(int, input().split())
# 국가간의 인구 수 분포
arr = []
for i in range(N):
    arr.append(list(map(int, input().split())))

# 병합 가능한 국가의 모든 국경선을 열어두어야 한다.
queue = deque()
visited = [[0] * N for _ in range(N)]
beonghaped = [] # 병합 대상이 되는 좌표들을 넣어줄 것

# 덩어리를 바꿔주는 역할
flag = 1
ans = 0
while flag:
    flag = 0
    ans += 1
    for i in range(N):
        for j in range(N):
            if visited[i][j] == 0: # 아직 탐색하지 않은 덩어리라면
                queue.append([i, j, arr[i][j]])
                # bfs() # 지금 상태면 병합 가능한 애들을 모두 한번에 합쳐버림
                # 이거 한번 돌때마다 한 덩어리에 대한 병합 가능 여부가 판단됨. 그 때 병합 가능한 칸의 정보를 받고, 그 칸들의 값의 합과 칸의 수를 받아 사용해야 할 것 같다.
                count, plus = bfs(i, j)
                for k in beonghaped:
                    arr[k[0]][k[1]] = plus // count

                beonghaped = []
    visited = [[0] * N for _ in range(N)]

print(ans - 1) # 최초 실행시 병합 요소가 없더라도 한번은 더해지기 때문에, 1을 빼주어야 한다.

그 결과 성공판정을 받았다. 어려운 문젠 아니었지만 막히던 부분을 스스로 해결하고 나니 뭔가 짜릿한 맛이 있는 것 같다.