Swift) 프로그래머스 - 전력망을 둘로 나누기 - lv2

https://school.programmers.co.kr/learn/courses/30/lessons/86971#

 

프로그래머스

코드 중심의 개발자 채용. 스택 기반의 포지션 매칭. 프로그래머스의 개발자 맞춤형 프로필을 등록하고, 나와 기술 궁합이 잘 맞는 기업들을 매칭 받으세요.

programmers.co.kr

문제를 간단히 요약하면 서로 연결되어 있는 집합끼리의 개수의 차이를 구하는 문제. 보자마나 Union-Find 알고리즘을 이용하면 풀 수 있겠구나 생각이 들었다. 와이어를 하나씩 자른 모든 경우의 수에 각 그룹 송전탑의 수 차이를 완전탐색해서 풀었다. Union-Find 알고리즘을 구현할 수 있다면 막힘없이 풀 수 있는 문제라고 생각

import Foundation

func calculateTwoGroupNum(_ n: Int, _ wires: [[Int]]) -> Int {
    var parent = [Int](0...n)
    func getParent(_ x: Int) -> Int {
        if parent[x] == x {
            return x
        }
        parent[x] = getParent(parent[x])
        return parent[x]
    }
    
    func unionParent(_ a: Int, _ b: Int) {
        let aParent = getParent(a) //1
        let bParent = getParent(b) // 2
        if (aParent < bParent) {
            parent[bParent] = aParent
        } else {
            parent[aParent] = bParent
        }
    }
    
    for wire in wires {
        unionParent(wire[0], wire[1])
    }
    
    var dict: [Int:Int] = [:]
    for i in 1..<parent.count {
        let rootParent = getParent(i)
        if dict[rootParent] != nil {
            dict[rootParent]! += 1
        } else {
            dict[rootParent] = 1
        }
    }
    let countArr = Array(dict.values)
    return abs(countArr[0] - countArr[1])
}

func solution(_ n:Int, _ wires:[[Int]]) -> Int {
    var answer = Int.max
    for num in 0..<wires.count {
        var copyWires = wires
        copyWires.remove(at: num)
        copyWires.sort { first, second in
            if first[0] == second[0] {
                return first[1] < second[1]
            }
            return first[0] < second[0]
        }
        answer = min(calculateTwoGroupNum(n, copyWires), answer)
    }
    
    return answer
}