링크

https://www.acmicpc.net/problem/16562

풀이

친구의 친구도 친구다 ⇒ MST 문제

1. 문제에서 주어진 정보를 통해 Union-find 자료구조를 완성한다.
2. 가중치가 적은 순으로 친구에 추가하면서 MST를완성한다.
3. MST가 완성되지 않은 경우 “Oh no” 출력
4. MST가 완성됬지만 내가가진 돈보다 많이드는경우 “Oh no” 출력
5. MST완성 비용 가진돈보다 작거나 같은경우 비용 출력

코드

//
//  main.swift
//  swiftPractice
//
//  Created by yangjs on 2021/12/25.
//

import Foundation
func union(_ a:Int,_ b:Int){
    let a = find(a)
    let b = find(b)
    if a>b{
        parent[b]=a
        return
    }
    parent[a]=b
}
func find(_ a:Int)->Int{
    if parent[a]==a{
        return a
    }
    let tmp = find(parent[a])
    parent[a]=tmp
    return parent[a]
}
var NMK = readLine()!.split(separator: " ").map{Int($0)!}
let cost = readLine()!.split(separator: " ").map{Int($0)!}
//[가중치,노드번호] 로 배열 재구성 => 가중치 순으로 정렬하기 위해
var info = (0..<cost.count).map{ i->Array in return [cost[i],i] }
info.sort {$0[0]<$1[0]}

var parent:[Int] = (0..<NMK[0]).map{$0}

//문제에서 주어진 정보로 유니온파인드 구성
for _ in 0..<NMK[1]{
    let ab = readLine()!.split(separator: " ").map{Int($0)!}
    if find(ab[0]-1) != find(ab[1]-1){
        union(ab[0]-1,ab[1]-1)
    }
}
parent.append(NMK[0])
var linkCost = 0
//크루스칼 알고리즘 
for costF in info{
    if find(NMK[0]) != find(costF[1]){
        linkCost+=costF[0]
        union(NMK[0],costF[1])
    }
}
var fin = true
//MST가 완성된 경우
for i in 0..<parent.count-1{
    if find(parent[i]) != find(parent[i+1]){
        print("Oh no")
        exit(0)
    }
}
//완성됬지만 예산 초과인경우, 완성된경우
linkCost > NMK[2] ? print("Oh no"):print(linkCost)