문제 링크: https://www.acmicpc.net/problem/1647

Thumbnail: Photo by Matthieu Rochette(Unsplash)

최소 신장 트리(MST)를 구하는 문제입니다.

풀이

"하나의 도시를 둘로 나누고, 각 도시를 연결하는 도로의 가중치 합이 최소가 되게 하라"

문제에서 요구하는 사항은 위와 같습니다. 이를 간단히 생각해보면, 전체 도시의 최소 신장 트리(Minimum Spanning Tree; MST)를 구한 뒤, 제일 가중치가 큰 간선 하나만 제거하면 두 도시를 최소 가중치 합으로 연결할 수 있습니다.

그렇다면, 전체 도시의 MST를 구하기 위해서는 어떤 알고리즘을 써야 할까요? 대표적인 방법으로는 Kruskal 알고리즘과 Prim 알고리즘이 있습니다. 이번 문제에는 Kruskal 알고리즘을 사용했습니다.

Union-Find을 활용해 사이클 여부 파악하기

Kruskal 알고리즘을 구현하기에 앞서, Union-Find 알고리즘에 대해 이해해야 합니다.

MST를 만드는 과정에서, 만약 잘못된 간선을 선택해 사이클이 생긴다면 이는 더이상 신장 트리가 아니게 됩니다. 따라서, 새로운 간선을 MST에 추가할 때 이 간선이 사이클을 형성하는지 미리 파악해야 합니다. 이를 위해 Union-Find 알고리즘을 사용할 수 있습니다.

Union-Find 알고리즘은 분리 집합의 개념을 구현하는 방법으로, 각 그룹을 트리로 표현하며 임의의 두 원소가 서로 공통된 조상을 갖는지 판단하거나, 두 트리를 합치는 작업을 수행합니다.

분리 집합이란?

전체 집합 \(U\)에 대해, 서로 겹치는 원소를 가지지 않는 두 부분 집합 \(A, B\)를 말합니다. 조건은 다음과 같습니다.

• \(A \subset U\), \(B \subset U\)
• \(A \cap B = \O\)
• \(A \cup B = U\)

꼭 두 집합이 아니어도, 전체 집합 \(U\)의 원소를 적절히 쪼개어 가지는 임의의 부분집합들을 분리 집합이라고 할 수 있습니다.

Union-Find 알고리즘은 다음 두 가지 함수를 사용합니다.

• find_root(x)
  • \(x\)가 속한 분리 집합의 최상위 노드를 반환합니다.
• union_root(x, y)
  • \(x\)가 속한 분리 집합과 \(y\)가 속한 분리 집합을 병합합니다.

처음 전체 집합의 원소들은 자기 자신을 최상위 노드로 가집니다. 이후, union_root를 통해 집합들을 병합합니다.

# Union Find 알고리즘 구현
parent = [i for i in range(N + 1)]
def find_root(x):
    """x가 속한 분리 집합의 최상위 노드를 반환한다.
    분리 집합(트리)의 깊이를 줄이기 위해, 임의의 노드 x에 대해 최상위 노드를 탐색한 뒤 최상위 노드를 바로 부모 노드로 설정한다.
    """
    if parent[x] == x:
        return x
    parent[x] = find_root(parent[x])
    return parent[x]

def union_root(x, y):
    """x와 y가 속한 두 분리 집합을 병합한다."""
    root_x = find_root(x)
    root_y = find_root(y)

    if root_x != root_y:
        parent[root_y] = root_x

Kruskal 알고리즘의 구현

Kruskal 알고리즘은 간선을 중심으로 MST를 구축합니다. 따라서, 사전에 간선이 가중치에 따라 정렬되어 있어야 합니다. 이 문제에서는 간선 정보가 정점1 정점2 가중치 순으로 주어지므로, 다음과 같은 코드로 정렬했습니다.

edges = sorted((tuple(map(int, input().split())) for _ in range(M)), key=lambda edge: edge[2]) # 전체 간선을 입력받은 뒤 가중치 기준으로 정렬한다.

이후, 정렬된 각 간선에 대해 Union-Find 알고리즘으로 사이클이 형성되는지 판단하고, 사이클이 형성되지 않는 간선은 MST에 추가해줍니다.

\(N\)개의 정점을 가지는 그래프에 대해, MST에 \(N-1\)개의 간선이 추가될 때 까지 위 과정을 반복해주면 됩니다.

# Kruskal MST 알고리즘 구현
edges = sorted((tuple(map(int, input().split())) for _ in range(M)), key=lambda edge: edge[2]) # 전체 간선을 입력받은 뒤 가중치 기준으로 정렬한다.

# Kruskal 알고리즘
mst = deque()
total_cost = 0
for i in range(M):
    x, y, c = edges[i]

    if find_root(x) == find_root(y):
        continue    # 사이클이 발생하는 간선은 MST에 넣지 않는다.
        
    mst.append((x, y, c))
    total_cost += c
    union_root(x, y)

    if len(mst) == N - 1: # N개의 정점에 대해 N-1개의 간선을 찾았다면 MST를 완성한 것이니 탐색을 종료한다.
        break

정답 구하기

문제에서 원하는 정답은 하나의 도시의 MST가 아닌, 두개의 도시의 MST의 최소 합입니다.

전체 집을 연결하는 MST에서 가장 가중치가 큰 간선 하나만 제거하면, 남는 2개의 신장 트리는 각각의 도시를 연결하는 최소 신장 트리가 됩니다.

print(total_cost - mst[N-2][2])

전체 코드

from collections import deque

input = open(0).readline
N, M = map(int, input().split())

# Union Find 알고리즘 구현
parent = [i for i in range(N + 1)]
def find_root(x):
    """x가 속한 분리 집합의 최상위 노드를 반환한다.
    분리 집합(트리)의 깊이를 줄이기 위해, 임의의 노드 x에 대해 최상위 노드를 탐색한 뒤 최상위 노드를 바로 부모 노드로 설정한다.
    """
    if parent[x] == x:
        return x
    parent[x] = find_root(parent[x])
    return parent[x]

def union_root(x, y):
    """x와 y가 속한 두 분리 집합을 병합한다."""
    root_x = find_root(x)
    root_y = find_root(y)

    if root_x != root_y:
        parent[root_y] = root_x

# Kruskal MST 알고리즘 구현
edges = sorted((tuple(map(int, input().split())) for _ in range(M)), key=lambda edge: edge[2]) # 전체 간선을 입력받은 뒤 가중치 기준으로 정렬한다.

# Kruskal 알고리즘
mst = deque()
total_cost = 0
for i in range(M):
    x, y, c = edges[i]

    if find_root(x) == find_root(y):
        continue    # 사이클이 발생하는 간선은 MST에 넣지 않는다.
        
    mst.append((x, y, c))
    total_cost += c
    union_root(x, y)

    if len(mst) == N - 1: # N개의 정점에 대해 N-1개의 간선을 찾았다면 MST를 완성한 것이니 탐색을 종료한다.
        break

# 최소 비용으로 두개의 도시를 만들기 위해서는, 현재 도시를 구성하는 MST를 만든 뒤 가장 비용이 높은 간선을 제거하면 된다.
print(total_cost - mst[N-2][2])