1647. 도시 분할 계획 (Python)

1647번: 도시 분할 계획

 

풀이

 최소 신장 트리(MST, Minimum Spanning Tree) 개념을 사용하여 해결할 수 있다. MST를 만들기 위한 알고리즘 중 크루스칼 알고리즘을 사용하였다. 유지비를 오름차순으로 정렬하여 작은것부터 선택하여 사이클이 발생하지 않도록 그래프를 만들어나가면 된다. 

 

 원래 N개의 정점이 있는 그래프를 MST로 만들기 위해서는 N-1개의 간선이 필요한데 문제에서는 마을을 두 개의 분리된 마을로 분할한다고 하였다. 그렇기때문에 N-2개의 간선만 고른다면 두 개의 분리된 마을로 만들 수 있다.

 

코드

import sys 

input = sys.stdin.readline

# 현재 집의 루트 노드 탐색
def find(v):
    if p[v] == v:
        return v
    p[v] = find(p[v])
    return p[v]


# 두 집의 관계 설정
def union(v1, v2):
    v1 = find(v1)
    v2 = find(v2)

    if v1 == v2:
        return

    p[v2] = v1


if __name__ == '__main__':
    N, M = map(int, input().split()) # N: 집의 개수, M: 길의 개수 
    graph = [] # 집들을 연결하는 길 
    p = [x for x in range(N+1)] # 각 집의 루트 노드 

    for i in range(M):
        a, b, c = map(int, input().split())
        graph.append((a, b, c))
    graph.sort(key=lambda x: x[2]) # 길을 유지비를 기준으로 오름차순 정렬 

    edgeCount = 0 # 길의 개수  
    cost = 0 # 유지비의 합 

    for (a, b, c) in graph:
        if edgeCount == N-2: # N-2개를 선택했다면 종료 
            break
        
        # 현재 길을 선택했을 때 사이클이 발생한다면 선택하지 않음 
        if find(a) == find(b): 
            continue

        union(a, b) # a집과 b집을 연결
        edgeCount += 1 # 길의 개수 1 증가 
        cost += c # 유지비만큼 합 증가 

    print(cost)