[백준, python] 14621번 - 나만 안되는 연애
백준 문제풀이 시작
문제
깽미는 24살 모태솔로이다.
깽미는 대마법사가 될 순 없다며 자신의 프로그래밍 능력을 이용하여 미팅 어플리케이션을 만들기로 결심했다.
미팅 앱은 대학생을 타겟으로 만들어졌으며 대학교간의 도로 데이터를 수집하여 만들었다.
이 앱은 사용자들을 위해 사심 경로를 제공한다. 이 경로는 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을 출력한다.)
풀이
풀이 확인하기
이번 문제는 1197번 문제와 거의 동일하다. 해당 링크.를 참조하자. 여기서 다른 점은 edge 추가 과정에서 edge가 연결하는 두 vertex가 둘 다 M이거나 W일 경우 추가하지 않고, 또한 MST를 만든 뒤 edge의 수가 vertex의 수 - 1가 아닐 경우 모든 학교가 연결되어있지 않다는 뜻이므로 -1을 출력한다.
import sys
input = sys.stdin.readline
def find(p, a):
if(p[a] != a):
p[a] = find(p, p[a])
return p[a]
def union(p, a, b):
a = find(p, a)
b = find(p, b)
if(a < b):
p[b] = a
else:
p[a] = b
v, e = map(int, input().split())
p = [0] * (v + 1)
edge = []
cost = 0
s = 0
for i in range(1, v+1):
p[i] = i
lis = input().strip().split()
for _ in range(e):
a, b, c = map(int, input().split())
if(lis[a-1] != lis[b-1]):
edge.append((c, a, b))
edge.sort()
for ed in edge:
c, a, b = ed
if(find(p, a) != find(p, b)):
union(p, a, b)
cost += c
s+=1
if(s != v-1):
print(-1)
else:
print(cost)