[Baekjoon 14621] 나만 안되는 연애

문제

💛Ⅲ [Baekjoon 14621 - 나만 안되는 연애]

문제 풀이 방법

모든 노드를 최소 거리의 비용으로 연결해야 하므로 최소 스패닝 트리 알고리즘인 크루스칼 알고리즘을 이용하여 구할 수 있다!

1. 성별이 다른 학교를 연결한 간선을 비용에 따라 오름차순으로 정렬한다.
2. 간선을 하나씩 확인하며 현재의 간선이 사이클을 발생시키는지 확인한다.
   1. 사이클이 발생하지 않는 경우, 최소 신장 트리에 포함시키며 그 때의 비용을 계산한다.
   2. 사이클을 발생하는 경우, 최소 신장 트리에 포함시키지 않는다.
3. 모든 간선에 대하여 2번 과정을 반복한다.

풀이 코드

Python

import sys
input = sys.stdin.readline

def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]

def union_parent(parent, a, b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

N, M = map(int, input().split())
schools = input().split()
edges = []
for _ in range(M):
    a,b,cost = map(int, input().split())
    if schools[a-1] != schools[b-1]:
        edges.append([a,b,cost])

parent = [i for i in range(N+1)]

edges.sort(key=lambda x:x[2])
result = 0
cnt = 1
for edge in edges:
    a,b,cost = edge
    if find_parent(parent, a) != find_parent(parent, b):
        union_parent(parent,a,b)
        result += cost
        cnt += 1
if cnt!=N:
    print(-1)
else:
    print(result)