[백준] 16234 - 인구 이동 / Python / 골드 5

🚩문제 주소

📄풀이

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 not visited[nx][ny]:
        diff = abs(country[x][y] - country[nx][ny])
        if l<=diff<=r:
          visited[nx][ny] = True
          q.append([nx,ny])
          temp.append([nx,ny])
          people += country[nx][ny]
          count += 1
  combine = people//count
  if count>1:
    check = True
    for x,y in temp:
      country[x][y] = combine

BFS를 활용하여 두 나라의 인구 수 차이가 L과 R 사이라면 방문 체크를 합니다.
그리고 새로운 배열에 방문한 나라의 좌표를 넣고, people(인구 수)에 더해주고, count(나라의 수)를 증가시킵니다.

count가 1 초과라면 인구 이동이 일어난 것이므로 check(인구 이동 체크 변수)를 True로 바꿔주고 인구 이동이 일어난 나라의 인구 수를 바꿉니다.

while True:
  check = False
  visited = [[False]*n for _ in range(n)]
  for i in range(n):
    for j in range(n):
      if not visited[i][j]:
        bfs(i,j,country,visited)
  
  if check:
    answer += 1
  else:
    break

check가 True라면 인구 이동이 일어난 것이므로 정답 +1을 하고, 아니라면 종료합니다.

PyPy3로 제출해야 통과합니다..

👨‍💻코드

from collections import deque

n,l,r = map(int, input().split())
country = []
dx, dy = [1,-1,0,0], [0,0,-1,1]
answer = 0
for _ in range(n):
  country.append(list(map(int, input().split())))

def bfs(a, b, country, visited):
  global check
  q = deque()
  temp = []
  q.append([a,b])
  temp.append([a,b])
  visited[a][b] = True
  
  people = country[a][b]
  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 not visited[nx][ny]:
        diff = abs(country[x][y] - country[nx][ny])
        if l<=diff<=r:
          visited[nx][ny] = True
          q.append([nx,ny])
          temp.append([nx,ny])
          people += country[nx][ny]
          count += 1
  
  combine = people//count
  if count>1:
    check = True
    for x,y in temp:
      country[x][y] = combine

while True:
  check = False
  visited = [[False]*n for _ in range(n)]
  for i in range(n):
    for j in range(n):
      if not visited[i][j]:
        bfs(i,j,country,visited)
  
  if check:
    answer += 1
  else:
    break

print(answer)