[이것이 코딩 테스트다] 도시 분할 계획 with Python

🔒 문제

 동물원에서 막 탈출한 원숭이 한 마리가 세상 구경을 하고 있다. 어느 날 원숭이는 '평화로운 마을'에 잠시 머물렀는데 마침 마을 사람들은 도로 공사 문제로 머리를 맞대고 회의 중이었다.
 마을은 N개의 집과 그 집들을 연결하는 M개의 길로 이루어져 있다. 길은 어느 방향으로든지 다닐 수 있는 편리한 길이다. 그리고 길마다 길을 유지하는데 드는 유지비가 있다.
 마을의 이장은 마을을 2개의 분리된 마을로 분할할 계획을 세우고 있다. 마을을 분할할 때는 각 분리된 마을 안에 집들이 서로 연결되도록 분할해야 한다. 각 분리된 마을 안에 있는 임의의 두 집 사이에 경로가 항상 존재해야 한다는 뜻이다. 마을에는 집이 하나 이상 있어야 한다.
그렇게 마을의 이장은 계획을 세우다가 마을 안에 길이 너무 많다는 생각을 하게 되었다. 일단 분리된 두 마을 사이에 있는 길들은 필요가 없으므로 없앨 수 있다. 그리고 각 분리된 마을 안에서도 임의의 두 집 사이에 경로가 항상 존재하게 하면서 길을 더 없앨 수 있다. 마을의 이장은 위 조건을 만족하도록 길들을 모두 없애고 나머지 길의 유지비의 합을 최소로 하고 싶다. 이것을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 집의 개수 N, 길의 개수 M이 주어진다.
N은 2 이상 100,000 이하인 정수이고, M은 1 이상 1,000,000 이하인 정수이다.
그 다음 줄부터 M 줄에 걸쳐 길의 정보가 A, B, C 3개의 정수로 공백으로 구분되어 주어지는데 A번 집과 B번 집을 연결하는 길의 유지비가 C(1 <= C <= 1,000) 라는 뜻이다.

7 12
1 2 3
1 3 2
3 2 1
2 5 2
3 4 4
7 3 6
5 1 5
1 6 2
6 4 1
6 5 3
4 5 3
6 7 4

출력

8

✅ 해설

• 한 마을 내 임의의 두 집 사이에는 경로가 항상 존재하며, 유지비의 합을 최소로 하기 위해선
  • 마을 사이에 사이클이 발생해서는 안 됨
  • 마을 사이에 사이클이 발생하는 순간 낭비되는 유지비가 생기기 때문
• 따라서 사이클이 없는 트리 구조를 계산하는 크루스칼 알고리즘 사용
  • 두 노드와 유지비를 우선순위 큐에 담는다.
    • 우선순위 큐는 유지비가 작은 순으로 자동 정렬해준다.
  • 유지비가 작은 순으로 그래프를 형성하되, 사이클이 생긴다면 건너뛴다.
  • 위 과정을 반복한다.

    # 마을을 두 개로 나누어야 하므로
    # 간선의 개수가 (n-2)개가 될 때까지만 반복
    while queue and road < n-2:
        # 간선의 개수 + 1
        road += 1
        # 우선순위 큐에서 정보 가져오기
        cost, n1, n2 = heapq.heappop(queue)

        # 사이클 판별
        if find_top(n1) != find_top(n2):
            union(n1, n2)
            result += cost

• 이때, 마을이 두 개로 나누어져야 하므로 길 하나를 삭제한다.
  • 마을이 하나이고 집의 개수가 N 개일 때, 최소 길의 수는 N-1 개이고
  • 마을이 두 개이고 집의 개수가 N 개일 때, 최소 길의 수는 N-2 개이고
  • ...
  • 마을이 M 개이고 집의 개수가 N 개일 때, 최소 길의 수는 N-M 개이이다.
  • 따라서, 길의 수가 N-2 개일 때까지만 반복한다.

💻 코드

import heapq

# 집의 개수 n과 길의 개수 m 입력받기
n, m = map(int, input().split())

# 우선순위 큐로 사용할 배열
queue = []

# 길 입력받기
graph = [[] for _ in range(n+1)]
for _ in range(m):
    n1, n2, cost = map(int, input().split())
    # 우선순위 큐에 삽입
    heapq.heappush(queue, (cost, n1, n2))

# 부모 노드를 저장하는 배열 & 초기화
parent = []
for i in range(n+1):
    parent.append(i)

# 최상위 부모 노드를 찾는 함수
def find_top(n):
    if n != parent[n]:
        parent[n] = find_top(parent[n])
    return n

# 합치기 함수
def union(n1, n2):
    p_n1 = find_top(n1)
    p_n2 = find_top(n2)
    if p_n1 < p_n2:
        parent[p_n2] = p_n1
    else:
        parent[p_n1] = p_n2

# 크루스칼 알고리즘
def kruskal():
    # 총 유지비
    result = 0
    # 간선의 개수
    road = 0

    # 마을을 두 개로 나누어야 하므로
    # 간선의 개수가 (n-2)개가 될 때까지만 반복
    while queue and road < n-2:
        # 간선의 개수 + 1
        road += 1
        # 우선순위 큐에서 정보 가져오기
        cost, n1, n2 = heapq.heappop(queue)

        # 사이클 판별
        if find_top(n1) != find_top(n2):
            union(n1, n2)
            result += cost

    return result

print(kruskal())