[백준-파이썬] 14621: 나만 안되는 연애

문제로 가기

싸피에서 prim 알고리즘을 배웠다.!
관련된 문제를 풀고 싶어서 푼 문제!
Prim에서 조금 응용해야 한다.
여기에서는 남초 대학, 여초 대학을 번갈아서 연결하는 조건까지 추가되었다!
힙에 넣을 때, 이전 노드의 타입(여초인지 남초인지) & 현재 노드의 타입까지 추가해서 풀었다 !

import heapq
import sys

input = sys.stdin.readline

N, M = map(int, input().split())
university_type = [0] + input().split()
visited = [False] * (N + 1)
distance = [9999999999] * (N + 1)

adj = [[] for _ in range(N + 1)]

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


def prim():
    Q = []
    ans = 0
    cnt = 0  # 연결한 학교 수
    # 출발점 초기화
    heapq.heappush(Q, (0, 1, 'no' , university_type[1]))  # 거리, 위치, 연결된 노드의 타입, 타입

    while Q:
        curr_dis, curr_idx, adj_idx_type, type = heapq.heappop(Q)
        if visited[curr_idx]:
            continue

        if adj_idx_type == type:
            continue

        now_type = type
        visited[curr_idx] = True
        ans += curr_dis
        cnt += 1
        if cnt == N:
            return ans

        for node, weight in adj[curr_idx]:
            if not visited[node]:  # 인접, 방문 체크
                if university_type[node] != now_type:
                    heapq.heappush(Q, (weight, node, university_type[curr_idx], university_type[node]))

    return -1


print(prim())