[백준] 14621번 나만 안되는 연애 파이썬 코드

코드

# 모든 대학교로 이동 가능해야한다 -> 신장 트리
# 최소 비용 -> 최소 신장 트리
# 크루스칼 알고리즘 이용

import sys
input = sys.stdin.readline

def find(a, lst):
  if a != lst[a]:
    lst[a] = find(lst[a],lst)
  return lst[a]

def union(a,b,lst):
  a = find(a, lst)
  b = find(b, lst)
  if a < b:
    lst[b] = a
  else:
    lst[a] = b

n, m = map(int, input().split())
school = ['x', *input().rstrip().split()] # 인덱싱을 쉽게 하기 위해 'x'추가
parent = [i for i in range(n+1)]
edge = []

for _ in range(m):
  a,b,c = map(int, input().split())
  edge.append((c,a,b))

edge.sort(key=lambda x: x[0]) # 비용을 기준으로 오름차순 정렬

count = 0 # 연결된 개수를 알기 위해 추가
res = 0 # 비용을 출력해줄 변수
for i in range(m):
  cost, v1, v2 = edge[i]
  if find(v1, parent) != find(v2, parent) and school[v1] != school[v2]:
    union(v1,v2, parent)
    res += cost
    count += 1
print(res) if count == n-1 else print(-1) # 모두 연결이 되지 않으면 -1