[백준] 14621 나만 안되는 연애 (Python/파이썬)

<문제 링크>

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

14621번: 나만 안되는 연애

<풀이 전략>

1. 기본적인 MST 문제이다.

2. 대학교 남초/여초 여부를 딕셔너리 형태로 저장해서 크루스칼 알고리즘에서 다음 간선 선택 시 다른 종류의 대학교일 경우만 선택하도록 했다.

3. 그 외에는 조건을 만족하는 mst가 없을 경우 -1을 출력하도록 해야된다.

 

<정답 코드>

import sys, os, io, atexit

input = lambda: sys.stdin.readline().rstrip('\r\n')
stdout = io.BytesIO()
sys.stdout.write = lambda s: stdout.write(s.encode("ascii"))
atexit.register(lambda: os.write(1, stdout.getvalue()))

# 220927 14621 나만 안되는 연애

# 정답코드


# 최소 신장트리를 위한 유니온파인드 함수
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]


def merge(x, y):
    x = find(x)
    y = find(y)

    if x < y:
        parent[y] = x
    else:
        parent[x] = y
    
    return


# N: 학교의 개수, M: 도로의 개수 
N, M = map(int, input().split())
# univ: 각 대학교의 남초/여초 여부
univ = dict(enumerate(list(input().split()), start=1))

# parent: 부모 노드 배열
parent = [i for i in range(N + 1)]

# graph: 대학교 그래프 연결 정보
graph = []
for _ in range(M):
    node1, node2, weight = map(int, input().split())
    graph.append((weight, node1, node2))

graph.sort()

# 최소 신장 트리로 문제 해결
def solve():

    # ans: 경로의 최소 길이, cnt: 경로로 연결된 대학교 개수
    ans = 0
    cnt = 0

    for i in range(M):
        distance, node1, node2 = graph[i]
        if  univ[node1] != univ[node2]:
            if find(node1) != find(node2):
                # 대학교 연결
                merge(node1, node2)

                # 경로 길이 계산 후, 대학교 세주기
                ans += distance
                cnt += 1

                # 모든 대학교가 연결되었을 경우 정답 출력 후 반환
                if cnt == N - 1:
                    print(ans)
                    return

    # 모든 학교를 연결하는 경우가 없을 경우 -1 출력 후 반환
    print(-1)
    return

solve()