문제링크

16234번: 인구 이동

문제설명

NN크기의 땅, 11마다 나라 존재 → r행 c열 나라의 인구는 A[r][c]명

더이상 인구이동이 없을때까지 인구이동

1. 국경열기 : 두 나라간 인구차이가 L명이상 R명 이하

2. 연합 만들기 : 국경 열린 국가끼리

   → 연합에 속한 각 국가들의 인구 = (연합 총 인구수)/(연합 이루는 국가 수)

3. 국경 선 닫기
   
   

문제 풀이

한 turn에 모든 나라를 돌면서 연합을 만들고, 한번에 이동시키는 것이 1번의 인구이동이다.
-> bfs() 함수에서 연합만들기+해당 연합에서 인구이동 까지 구현했을때,

3 5 15
10 15 20
20 30 25
40 22 5

와 같은 테스트 케이스에서
첫번째 인구이동
bfs(0,0)실행시

27 27 27
27 27 27
40 27 5

와 같은 결과를 얻게 된다.
bfs(2,2)를 실행했을때 "27"은 현재 turn의 인구이동의 결과임으로 이번 차례에 (2,2)에서 27로는 union을 만들 수 없다는 점을 주의해야 한다.
따라서 인구이동을 반복하는 while루프 안에 각 나라의 연합여부를 체크하는 united 배열을 놓고, bfs에서 연합에 추가하려는 나라가 이미 이번 turn에 인구이동을 마쳤는지 체크해줘야 한다 !

시도 1

from collections import deque
import sys

input = sys.stdin.readline
n,l,r = map(int, input().split())

worldmap = []
for _ in range(n):
    worldmap.append(list(map(int, input().split())))

dx = [-1, 0, 1, 0]
dy = [0, -1, 0, 1]

result = 0

def process(x, y, index):
    #(x,y)와 연결된 연합 정보 담기
    united = []
    united.append((x,y))

    #BFS로 연결된 연합 확인
    q = deque()
    q.append((x,y))
    union[x][y] = index
    #연합의 인구수
    population = worldmap[x][y]
    #연합에 속한 나라 수
    count = 1
    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 union[nx][ny] == -1:
                if l<= abs(worldmap[x][y] - worldmap[nx][ny]) <=r:
                    q.append((nx,ny))
                    union[nx][ny] = index
                    population += worldmap[nx][ny]
                    count += 1
                    united.append((nx,ny))
    for i,j in united:
        worldmap[i][j] = population//count

total_count = 0

while True:
    union = [[-1]*n for _ in range(n)]
    index = 0
    for i in range(n):
        for j in range(n):
            if union[i][j] == -1:
                process(i,j, index)
                index += 1
    #모든 인구이동 후
    #print(union)
    if index == n*n:
        break
    total_count += 1

print(total_count)

왜 계속 시간초과가 날까 했는데 pypy3으로 제출하니까 됐음 ㅋ...

시도 2

두달이 지난 오늘(1/20) 다시 풀어봤다,, 인구이동을 모든 연합에 대해서 한번에 처리하는 부분에서 헤맸음 ^_ㅠ,,

import sys
from collections import deque

input = sys.stdin.readline

dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]

N, L, R = map(int, input().split())
countries = []
answer = 0
for _ in range(N):
    countries.append(list(map(int, input().split())))

def in_range(x, y):
    if x in range(N) and y in range(N):
        return True
    return False

def check_union(x, y, nx, ny):
    if abs(countries[x][y] - countries[nx][ny]) in range(L, R+1):
        return True
    return False

def bfs(x, y, index):
    queue = deque([(x,y)])
    union = [(x, y)]
    total = countries[x][y]

    while queue:
        x, y = queue.popleft()

        for i in range(4):
            nx, ny = x+dx[i], y+dy[i]
            if in_range(nx, ny) and check_union(x, y, nx, ny) and united[nx][ny] == -1:
                queue.append((nx, ny))
                union.append((nx, ny))
                united[nx][ny] = 1
                total += countries[nx][ny]
    
    population = total//len(union)
    for x, y in union:
        countries[x][y] = population

while True:
    united = [[-1]*N for _ in range(N)]
    index = 0
    for i in range(N):
        for j in range(N):
            #아직 합쳐지지 않은 나라에 대해서
            if united[i][j] == -1:
                united[i][j] = index
                bfs(i, j, index)
                index += 1
    if index == N*N:
        break
    answer += 1

print(answer)