[백준/BOJ] 14621번 - 나만 안되는 연애 (Python)

문제링크

https://www.acmicpc.net/problem/14621

14621번: 나만 안되는 연애

1. 핵심

• 알고리즘 유형: 최소 스패닝 트리
  • 크루스칼 알고리즘
• 성별이 다른, 남자대학교와 여자대학교만 연결되어야 한다.
  • 남-남, 여-여끼리 간선이 연결되면 안된다.

2. 코드(Python)

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

def union(parent, a, b):
    pa = find(parent, a)
    pb = find(parent, b)
    if pa == pb:
        return False
    parent[pa] = pb
    return True

def solution(n, adj):
	# 비용으로 정렬
    adj.sort(key=lambda x:x[4])
    parent =[i for i in range(n + 1)]
    sum_val = 0
    cnt = 0
    for gender_v, gender_u, v, u, cost in adj:
        if cnt == n - 1:
            break
        # 출발점과 도착점의 성별이 같으면 union 진행하지 않는다.
        if gender_v == gender_u:
            continue
        if union(parent, v, u):
            sum_val += cost
            cnt += 1
    if sum_val == 0 or cnt != n - 1:
        return -1
    return sum_val

n, m = map(int, input().split())
genders = [0] + list(input().split())
adj = []
for idx in range(m):
    v, u, cost = map(int, input().split())
    adj.append((genders[v], genders[u], v, u, cost))
    
print(solution(n, adj))