알고리즘 문제풀이

나만 안되는 연애(백준 14621)

wiojfe 2023. 7. 31. 15:32

  문제의 조건에 의하면 남학교와 여학교를 이어주는 길만 갖고 MST를 만들어야 한다. 

  그래서 입력으로 들어오는 모든 길을 저장하지 않고 조건을 만족하는 경우만 저장한다. 

import sys
import heapq 
input = sys.stdin.readline
n,m= map(int, input().split())
gender =list(map(str, input().split()))
gender = ["x"] + gender
gra= [[]for _ in range(n+1)]
for i in range(m):
    a,b,c= map(int, input().split())
    if (gender[a] == gender[b]):
        continue 
    else:
        gra[a].append((b,c))
        gra[b].append((a,c))
visited= [0 for _ in range(n+1)]
# print(visited)
ans =0 
def prim(start):
    global ans 
    q = [[0,start]]
    heapq.heapify(q)
    while q :
        cost, now = heapq.heappop(q)
        if(visited[now]==0):
            
            visited[now] = 1 
            ans += cost
            for nextnode,nextcost in gra[now]:
                if visited[nextnode] ==0:
                    heapq.heappush(q, (nextcost , nextnode))
                
                
        
prim(1)
if 0 in visited[1:]:
    print(-1)
    exit()
    
print(ans)