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

14621번: 나만 안되는 연애

 

문제

깽미는 24살 모태솔로이다. 깽미는 대마법사가 될 순 없다며 자신의 프로그래밍 능력을 이용하여 미팅 어플리케이션을 만들기로 결심했다. 미팅 앱은 대학생을 타겟으로 만들어졌으며 대학교간의 도로 데이터를 수집하여 만들었다.

이 앱은 사용자들을 위해 사심 경로를 제공한다. 이 경로는 3가지 특징을 가지고 있다.

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

만약 도로 데이터가 만약 왼쪽의 그림과 같다면, 오른쪽 그림의 보라색 선과 같이 경로를 구성하면 위의 3가지 조건을 만족하는 경로를 만들 수 있다.

이때, 주어지는 거리 데이터를 이용하여 사심 경로의 길이를 구해보자.

입력

입력의 첫째 줄에 학교의 수 N와 학교를 연결하는 도로의 개수 M이 주어진다. (2 ≤ N ≤ 1,000) (1 ≤ M ≤ 10,000)

둘째 줄에 각 학교가 남초 대학교라면 M, 여초 대학교라면 W이 주어진다.

다음 M개의 줄에 u v d가 주어지며 u학교와 v학교가 연결되어 있으며 이 거리는 d임을 나타낸다. (1 ≤ u, v ≤ N) , (1 ≤ d ≤ 1,000)

출력

깽미가 만든 앱의 경로 길이를 출력한다. (모든 학교를 연결하는 경로가 없을 경우 -1을 출력한다.)

코드

from heapq import heappop, heappush

n, m = map(int, input().split())
univ = ["N"] + input().split()
parent = [i for i in range(n+1)]
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    a = find(a)
    b = find(b)
    if a == b:  # 사이클을 형성
        return
    if a < b:
        parent[b] = a
    else:
        parent[a] = b
        
roads = []
for _ in range(m):
    u, v, d = map(int, input().split())
    heappush(roads, (d, u, v))

dist, cnt = 0, 1  
while roads:
    c, a, b = heappop(roads)
    if univ[a] == univ[b]:  # 같은 성별 학교
        continue
    if find(a) != find(b):  # 사이클 형성하는지 체크
        union(a,b)
        dist += c
        cnt += 1
    if cnt == n:  # 경로 완성
        print(dist)
        exit()
print(-1)

가장 적은 가중치를 갖는 스패닝트리를 찾는 문제이므로 크루스칼 알고리즘을 이용하여 풀었다.

다만 주의할 점은 같은 남초-남초 또는 여초-여초 학교끼리 이어주면 안되므로 univ 배열에 저장해놓은 번호별 학교의 성비 정보를 이용하여 같은 남초-여초 학교끼리만 이어지게 해주어야 한다.