Programming/Algorithm

[백준14621] / 나만 안되는 연애 - 최소 스패닝 트리(MST)

__narrrrrmm 2024. 12. 26. 23:21

 

백준14621

 

 

푼 날짜 : 2024.12.26

푼 문제 : [14621] / 나만 안되는 연애

사용한 언어 : python

알고리즘 : 최소 스패닝 트리(MST)

 

 

 

접근방식 :

일반적인 MST 알고리즘을 사용했다.

추가적으로 성별을 다르게 번갈아 가며 연결하기 위해 heap에 성별을 추가 요소로 적용했다.

 

문제의 조건은 다음과 같다.

  1. 사심 경로는 사용자들의 사심을 만족시키기 위해 남초 대학교와 여초 대학교들을 연결하는 도로로만 이루어져 있다.
  2. 사용자들이 다양한 사람과 미팅할 수 있도록 어떤 대학교에서든 모든 대학교로 이동이 가능한 경로이다.
  3. 시간을 낭비하지 않고 미팅할 수 있도록 이 경로의 길이는 최단 거리가 되어야 한다.

따라서 W면 M과 연결 M이라면 W와 연결하도록 구현했으며 이 과정에서 최단 거리를 계산할 수 있도록 구현했다.

 

 

 

 

코드 :

import sys
import heapq

input = sys.stdin.readline

N, M = map(int, input().split())
univ = list(input().split())

graph = [[] for _ in range(N+1)]
visited = [0 for _ in range(N+1)]

for _ in range(M):
    A, B, C= map(int, input().split())
    graph[A].append((B, C))
    graph[B].append((A, C))

# 비용, 시작지점, 시작지점의 성별
heap = [(0, 1, univ[0])]
sum_cost = 0

while heap:
    n_cost, n_edge, now_univ = heapq.heappop(heap)

    if not visited[n_edge]:
        sum_cost += n_cost
        visited[n_edge] = 1

        for e, c in graph[n_edge]:
            # 방문하지 않은 지점이며 현재의 성별과 다른 성별의 학교인 경우
            if not visited[e] and now_univ != univ[e-1]:
                heapq.heappush(heap, (c, e, univ[e-1]))

# 모두 방문했다면 비용 출력 아니라면 -1 출력
print(sum_cost if sum(visited[1:]) == N else -1)

 

 

 

골드3이라 살짝 어려운가 싶었지만 추가적인 조건하나만 더 구현해주면 됐던 문제!