[백준/Python] 16234 인구 이동

예제4, 5를 이해하는 데에 시간이 좀 오래 걸렸다..

아래는 이해한 예제4, 5를 정리해놓은 그림이다.

예제4

예제5

예제를 이해하고 코드를 작성하기 전에 했던 생각은,

 

1. 현재 국가(위치)를 기준으로 상하좌우 주변국들과의 인구수를 비교해야하니까 BFS를 사용하면 좋겠다

2. 그러면 인접한 나라들과의 인구수 비교를 해 연합국이 될 수 있는지를 확인할 수 있는 함수를 만들어야겠다

3. 연합국이 여러 개 일 수 있으니까, visited로 방문여부를 확인할 수 있어야겠다

4. 연합국들과의 인구조정이 된 후의 인구수로 상태를 계속 update해줘야겠다.

이정도였고, 코드를 작성하기 시작했다

from collections import deque

dx = [0, 0, 1, -1]  # 상하좌우
dy = [1, -1, 0, 0]  # 상하좌우


def bfs(i, j):
    q = deque()
    q.append((i, j))
    union = []
    union.append((i,j))
    while q:
        x, y = q.popleft()
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if 0 <= nx < N and 0 <= ny < N and visited[nx][ny] == 0:
                if L <= abs(A[x][y] - A[nx][ny]) <= R:  # 인구수 차이가 L이상 R이하일 때만 union에 append
                    visited[nx][ny] = 1
                    q.append((nx, ny))
                    union.append((nx, ny))

    return union


N, L, R = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(N)]
visited = [[0 for _ in range(N)] for _ in range(N)]
cnt = 0  # 몇 일이 지났는지 저장하는 변수

while True:
    visited = [[0 for _ in range(N)] for _ in range(N)]
    # flag는 방문하지 않은 나라를 다시 방문하기 전에 0으로 초기화 해주어서, 연합국이 없다면 바로 cnt가 출력되도록 함
    flag = 0
    for i in range(N):
        for j in range(N):
            if visited[i][j] == 0:
                visited[i][j] = 1
                countries = bfs(i, j)

                # 연합국이 있는 경우 flag = 1을 하고, 연합국들과의 인구조정이 된 후의 인구수로 바꿔 저장
                if len(countries) > 1:
                    flag = 1
                    population = sum(A[x][y] for x,y in countries) // len(countries)

                    for x, y in countries:
                        A[x][y] = population

    # 연합국이 더이상 없는 경우 cnt 출력
    if flag == 0:
        print(cnt)
        break
    cnt += 1  # 인구조정이 끝나면 하루가 지났다고 판단

풀이과정

1. 인접한 나라들과의 인구수 비교를 해 연합국이 될 수 있는지를 확인할 수 있는 함수 bfs를 만듦

    → 여느 평범한 BFS 알고리즘과 같이, 시작노드를 기준으로 상하좌우의 인접 노드를 확인하는데, while문을 사용해 큐에 저장된 인접노드가 없을 때까지 순차적으로 방문

2. flag = 0으로 시작해서, 연합국이 있는 경우에만 1로 바꿔주어 while문을 나갈 수 없게 함

3. 연합국이 있는 경우, 연합국들과의 인구조정(연합국 전체 인구/연합국 수) 인구수를 A에 업데이트 시켜주고

4. 다시 연합국이 될 수 있는지 확인하는 함수(bfs)를 적용해봄

5. 다시 bfs 함수에 넣기 전에 cnt를 1 증가시켜 하루가 지났음을 표시

6. 연합국이 없는 경우, flag가 0이므로 cnt를 출력하고 while문을 빠져나오고 끝남