[BOJ] [Python] 14621 : 나만 안되는 연애

모든 노드를 방문하고, 최단 거리를 구하기 위해 프림 알고리즘을 사용했다.

 

사심 경로의 첫 번째 특징은 남초 대학교와 여초 대학교들을 연결하는 도로로만 이루어져 있다는 것이다.

다시 말해서, 남초-남초 대학과 여초-여초 대학인 길은 선택해서는 안 된다.

 

이것은 현재 노드에서 현재 노드와 연결된 노드를 order 리스트에 넣어줄 때, 확인하면 된다.

 

 

import sys
import heapq

input = sys.stdin.readline

n, m = map(int, input().split())

graph = [[] for i in range(n + 1)]

gender = ['']
gender += list(map(str, input().split()))

for _ in range(m):
    u, v, d = map(int, input().split())
    graph[u].append((d, (u, v)))
    graph[v].append((d, (v, u)))

count = 0
dp = [0] * (n + 1)

order = []

now_node = 1 # 시작 노드
result = 0 # 가중치 합

# 초깃값
for edge in graph[1]:
    if gender[1] != gender[edge[1][1]]:
        heapq.heappush(order, edge)

while count < n-1:
    if not order:
        result = -1
        break

    node = heapq.heappop(order)

    A = node[1][0]
    B = node[1][1]
    val = node[0]

    if dp[A] == 0 or dp[B] == 0:
        dp[A] = 1
        dp[B] = 1

        result += val

        now_node = B
        count += 1

        for edge in graph[now_node]:
            if gender[now_node] != gender[edge[1][1]]:
                heapq.heappush(order, edge)
print(result)

 

 

모든 학교를 연결하는 경로가 없으면 -1을 출력한다.

 

예를 들어, 

3 3
m m m
1 2 3
2 3 4
3 1 5
이런 경우 -1을 출력해야 한다.