프로그래밍/Algorithm

백준 16234 인구 이동 파이썬

모지사바하 2021. 4. 1. 11:42

www.acmicpc.net/problem/16234

 

16234번: 인구 이동

N×N크기의 땅이 있고, 땅은 1×1개의 칸으로 나누어져 있다. 각각의 땅에는 나라가 하나씩 존재하며, r행 c열에 있는 나라에는 A[r][c]명이 살고 있다. 인접한 나라 사이에는 국경선이 존재한다. 모

www.acmicpc.net

 

와.. 이 문제 역대급이였다.

결국 혼자의 힘으로 풀어내긴 했지만 이틀 정도 헤맨것같다.

풀고나서 다른분들 풀이를 보니 내가 왜 이생각을 못했나.. 자괴감들고 괴로워.. ㅠ

import sys
from collections import deque

if __name__ == '__main__':
    direction = [(0, 1), (1, 0), (0, -1), (-1, 0)]

    n, left, right = map(int, sys.stdin.readline().rstrip().split())
    countries = [list(map(int, sys.stdin.readline().rstrip().split())) for _ in range(n)]


    # 인구 통합
    def population_integration(temp):
        for i in range(len(temp)):
            for j in range(len(temp[i])):
                if temp[i][j] > 0:
                    union_num = temp[i][j]
                    population = countries[i][j]
                    union_point = [(i, j)]
                    queue = deque([(i, j)])
                    temp[i][j] = 0
                    while queue:
                        v, h = queue.popleft()
                        for d in direction:
                            nv, nh = v + d[0], h + d[1]
                            if 0 <= nv < n and 0 <= nh < n and temp[nv][nh] and temp[nv][nh] == union_num:
                                temp[nv][nh] = 0
                                population += countries[nv][nh]
                                queue.append((nv, nh))
                                union_point.append((nv, nh))
                    union_population = int(population / len(union_point))
                    for a, b in union_point:
                        countries[a][b] = union_population


    # 국경 개방
    def open_national_line():
        integration_count = 0
        while True:
            temp = [[0] * n for _ in range(n)]
            union_flag = False

            union_num = 0
            visited = [[False] * n for _ in range(n)]
            for i in range(len(temp)):
                for j in range(len(temp[i])):
                    if temp[i][j] == 0:
                        union_num += 1
                        temp[i][j] = union_num
                        queue = deque([(i, j)])

                        while queue:
                            v, h = queue.popleft()
                            for d in direction:
                                nv, nh = v + d[0], h + d[1]
                                if 0 <= nv < n and 0 <= nh < n and not visited[nv][nh]:
                                    first = countries[v][h]
                                    second = countries[nv][nh]
                                    if left <= abs(first - second) <= right:
                                        visited[nv][nh] = True
                                        queue.append((nv, nh))
                                        union_flag = True
                                        temp[nv][nh] = union_num
                                        temp[v][h] = union_num

            if not union_flag:
                return integration_count

            integration_count += 1
            population_integration(temp)


    print(open_national_line())

# 2 10 20
# 0 30
# 50 10

 

개방할 수 있는 모든 국경을 개방해준다.

국경을 개방할때, 각 연합국의 번호를 임시 배열에 저장한다.

그 다음, 인구통합을 한다.

인구통합을 할때 임시 배열에 저장된 연합국의 번호를 보고 bfs 로 순회하며 같은 번호를 가진 연합국의 인구를 모두 합하면서

연합국의 각 국가 좌표를 저장한후, 인구를 재조정한다.

이 과정을 더 이상 연합국이 생기지 않을때까지 반복한다.

 

위는 내 풀이다. 그런데 저렇게 하게되면 효율이 좋지 않다.

개방할 수 있는 모든 국경을 개방한 후, 인구통합을 하면 전체 배열 순회를 2차례. BFS를 두차례 해야한다.

하지만 국경을 개방할때, 각 국가의 좌표와 인구를 저장하면 

하나의 연합국이 완료될때마다 바로 인구통합이 가능하므로 전체배열순회, BFS 를 한차례만으로 진행할 수 있다.

ㅠㅠ 다시 풀어봐야지.

아래는 모범답안.

from collections import deque
n, l, r = map(int, input().split())
arr = [list(map(int, input().split())) for _ in range(n)]
visit = [[False] * n for _ in range(n)]
q = deque()
dx = [0, 0, -1, 1]
dy = [1, -1, 0, 0]
ans = 0
while True:
  check = False
  for i in range(n):
    for j in range(n):
      if visit[i][j] == False:
        q.append((i,j))
        visit[i][j] = True
        s = arr[i][j]
        u_lst = [(i, j)]
        while q:
          x, y = q.popleft()
          for a in range(4):
            nx, ny = x + dx[a], y + dy[a]
            if 0 <= nx < n and 0 <= ny < n and visit[nx][ny] == False:
              if l <= abs(arr[nx][ny] - arr[x][y]) <= r:
                check = True
                visit[nx][ny] = True
                q.append((nx, ny))
                u_lst.append((nx, ny))
                s += arr[nx][ny]
        for x, y in u_lst:
          arr[x][y] = s // len(u_lst)
  if check:
    ans += 1
  else:
    break
  visit = visit = [[False] * n for _ in range(n)]
print(ans)