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

문제 출처:https://www.acmicpc.net/problem/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을 출력한다.)

예제 입력 1

5 7
M W W W M
1 2 12
1 3 10
4 2 5
5 2 5
2 5 10
3 4 3
5 4 7

예제 출력 1

34

---

풀이

 

문제에서의 세 가지 조건을 살펴보자.

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

 

1번 조건의 경우 남초 대학교(M)는 여초 대학교(W)와 연결되어 있고, 여초 대학교(W)는 남초 대학교(M)와 연결 되어 있다는 뜻이다. 즉, M과 M끼리는 직접 연결되어있지 않고, W와 W끼리는 직접 연결되어 있지 않다는 뜻이다.

 

2번 조건과 3번 조건의 경우, 모든 대학(M,W)이 연결되어 있으면서 경로의 길이가 최단 거리가 되어야 한다는 뜻이므로 최소 신장 트리를 구하라는 것과 같다.

 

이 조건들을 이용하여 코드를 작성하면 다음과 같다.

 

import sys
 
def find_parent(parent,x):
    if parent[x] != x:
        parent[x] = find_parent(parent,parent[x])
    return parent[x]
 
def union_parent(parent,a,b):
    a = find_parent(parent,a)
    b = find_parent(parent,b)
 
    if a<b:
        parent[b] = a
    else:
        parent[a] = b
 
 
n,m = map(int,sys.stdin.readline().split())
univ = ['']
univ.extend(sys.stdin.readline().split())
 
parent = [i for i in range(n+1)]
edges = []
for _ in range(m):
    u,v,d = map(int,sys.stdin.readline().split())
    if univ[u]!=univ[v]:
        edges.append((d,u,v))
 
edges.sort()
 
result = 0
l = 0
for d,u,v in edges:
    if find_parent(parent,u)!=find_parent(parent,v):
        union_parent(parent,u,v)
        result+=d
        l+=1
    if l==n-1:
        break
 
if l<n-1:
    result=-1
 
print(result)

• univ 리스트는 대학이 남초 대학교(M)인지 여초 대학교(W)인지를 저장할 리스트이다. 문제에서 학교의 번호가 1번부터 시작되므로 univ 리스트의 0번 index 자리에는 빈 문자열을 넣어주었다.
• 학교와 두 학교 사이의 거리를 입력받는다.
• u학교와 v학교가 둘 다 남초 대학교이거나 여초 대학교라면 어차피 이용할 수 없는 간선이므로 두 학교가 서로 다른 경우일 때만 간선 리스트에 추가한다. (u학교가 W이고 v학교가 M인 경우, u학교가 M이고 v학교가 W인 경우)
• 입력 받은 간선을 거리가 짧은 순으로 정렬하고, 짧은 간선부터 하나씩 살펴본다.
• 만약 연결한 간선이 n-1개라면, 최소 신장트리가 완성되었으므로 break를 한다.
• 만약 간선을 모두 살펴봤지만 연결한 간선이 n-1개보다 작다면 모든 학교를 연결하는 경로가 없으므로 result를 -1로 바꾼다.
• 최종적으로 result를 출력한다.