문제 설명
문제를 풀 당시에 골드2로 책정되어 있지만 사실 이보다 낮은 난이도 같습니다.
다른 최소 스패닝트리와 차이가 하나 있다면 이미 발전소와 연결되어 있는 도시는 굳이 한번 더 연결할 필요가 없기 때문에
연결하지 않으면 됩니다.
union 함수 부분에 로직을 몇가지 추가했는데, 발전소와 연결되어 있는지를 체크하기 위해서는 기존 방법(Index가 작은 곳으로 연결)대신
발전소를 기준으로 Parent 배열을 업데이트 해야 합니다.
그래서 발전소에 해당하는지를 체크하는 로직이 추가되었습니다.
소스 코드
Python
import sys
input = sys.stdin.readline
N,M,K = map(int,input().split()) #도시갯수, 케이블, 발전소
parent = [i for i in range(N+1)]
pp = list(map(int,input().split())) #power plant 발전소
arr = []
for i in range(M):
a,b,d = map(int,input().split())
arr.append([a,b,d])
arr = sorted(arr,key=lambda x:x[2]) #가치 오름차순으로 정렬
def find(x):
if parent[x]==x:return x
parent[x] = find(parent[x])
return parent[x]
def union(a,b):
a = find(a)
b = find(b)
if a==b:return False
app = a in pp
bpp = b in pp
if app and bpp:# 둘 다 발전소에 연결되어 있다면
return False#굳이 연결할 필요 없음
if app:#a가 발전소에 연결되어 있다면
parent[b]=a#발전소 쪽으로 연결
elif bpp:
parent[a]=b
else:#둘 다 발전소에 연결되어 있지 않다면?
if a<b:parent[b] = a
else:parent[a]=b
return True
answer = 0
for i in range(len(arr)):
a,b,d = arr[i]
if union(a,b):
answer+=d
print(answer)
Java8
#