[Algorithm/Swift] Union-Find 코드 개선

알고리즘 스터디를 하면서 프로그래머스의 섬 연결하기 문제로

Union-Find, 크루스컬 알고리즘을 접하게 됐고 우리는 이걸 짚고 넘어가자고 했다.

 

맨 처음 잘 모르는 채로 제출했던 코드는 아래 섬 연결하기 코드이다.

// [level 3] 섬 연결하기
import Foundation

var parents = [Int]()

func solution(_ n:Int, _ costs:[[Int]]) -> Int {
    
    parents = (0 ..< n).map{ $0 } // 각자 index의 값을 초기 값으로 설정
    let sortByCost =  costs.sorted { $0[2] < $1[2] } // cost 기준 오름차순 정렬
    var sum = 0 // 전체 합
    
    for bridge in sortByCost {
        // 사이클이 연결되어있지 않다면
        if !isCycle(left: bridge[0], right: bridge[1]) {
            // 다리 연결
            connectBridge(left: bridge[0], right: bridge[1])
            // 전체 합에 cost 누적
            sum += bridge[2]
        }
    }    
    return sum
}

// 사이클 형성 하는지
func isCycle(left:Int,right:Int) -> Bool {
    return parents[left] == parents[right] ? true : false
}

// 다리 연결
func connectBridge(left:Int,right:Int) {
    //oldParent로 right의 가장 부모를 가진 것들을 newParent를 left의 가장 부모로 바꿔준다.
    changeParent(oldParent: findParent(child: right), newParent: findParent(child: left))
}

// 부모 바꾸기
func changeParent(oldParent:Int,newParent:Int) {
    //oldParent의 값을 가진 index를 찾아내서 해당 index의 값을 newParent로 바꿔준다.
    parents.indices.filter {parents[$0] == oldParent}.forEach{parents[$0] = newParent}
}

//가장 위 부모 찾기
func findParent(child:Int) -> Int {
    //parrents의 child번째가 child라면 가장 위 부모 아니라면 재귀해줌
    return parents[child] == child ? child : findParent(child:parents[child])
}

 

그리고 Union-Find 알고리즘을 이해한 후 풀게된 집합의 표현 문제 풀이 코드는 이것이다.

// [Gold IV] 집합의 표현
let line = readLine()!.split(separator: " ").map{ Int(String($0))! }
var parents = Array(0...line[0])
//print(parents)

for _ in 0 ..< line[1] {
    let command = readLine()!.split(separator: " ").map{ Int(String($0))! }
    let a = command[1]
    let b = command[2]

    if command[0] == 0 { // Union
        union(a, b)
    } else { // command[0] == 1 -> Find
        print(find(a) == find(b) ? "YES" : "NO")
    }
}

func find(_ i: Int) -> Int {
    guard parents[i] != i else { return i }
        parents[i] = find(parents[i])
    return parents[i]
}

func union(_ a: Int, _ b: Int) {
    let a = find(a)
    let b = find(b)
    
    if a >= b {
        parents[a] = b
    } else {
        parents[b] = a
    }
}

 

여기서 skydreamer21님이 위에 섬 연결하기 코드가 아래 집합의 표현 문제 풀이보다 느린데

왜 그런지 이유를 찾아오기 미션을 줬다. 그래서 이것만큼은 제대로 파악해보고싶어서 구구절절 써봤다.

 

---

위 노트를 바탕으로 설명을 해보겠다!

섬 연결하기 코드부터 자세히 살펴보자면

parents = (0 ..< n).map{ $0 } // 각자 index의 값을 초기 값으로 설정
let sortByCost =  costs.sorted { $0[2] < $1[2] } // cost 기준 오름차순 정렬

여기서 이미 각 줄마다 .map 과 sorted 로 인해 O(n) 시간이 2번 소요된다.

for bridge in sortByCost {
    print("parents : \(parents)")
    // 사이클이 연결되어있지 않다면
    if !isCycle(left: bridge[0], right: bridge[1]) {
        // 다리 연결
        connectBridge(left: bridge[0], right: bridge[1])
        // 전체 합에 cost 누적
        sum += bridge[2]
        print("parents : \(parents)")
        print("left : \(bridge[0]), right: \(bridge[1])")
    }
}

주 코드 부분인 곳에서 parents 배열의 변화를 보면서 확인해보자

sortedByCost 의 0번째 인덱스부터 돌아보면 [0, 1, 1] 값을 가지고 있다.

bridge[0] = 0

bridge[1] = 1

cost = 1

 

현재 이런 간선들이 나열되어있지만 아직 연결된 건 없다고 가정했을 때

0 과 1은 아무것도 연결되어있지 않기 때문에 parents = [0, 1, 2, 3] 즉 자기 본인의 값을 부모로 갖고있다.

이 상태에서 사이클을 형성하는지 확인하는 함수는

// 사이클 형성 하는지
func isCycle(left:Int,right:Int) -> Bool {
    return parents[left] == parents[right] ? true : false
}

parents 배열에서 양쪽의 값이 같니? = 공통 부모이니? = 이미 연결되어있니? 를 확인한다.

근데 문제는 이 함수는 현재 값만 보고 있기 때문에 여기서 find를 쓴다면 아래 함수가 필요없다.

위 상황에서는 연결 안되어있으니 false 를 반환하고

connectBridge(left: bridge[0], right: bridge[1])

이 함수를 부른다.

그럼 이 함수는

// 다리 연결
func connectBridge(left:Int,right:Int) {
    //oldParent로 right의 가장 부모를 가진 것들을 newParent를 left의 가장 부모로 바꿔준다.
    changeParent(oldParent: findParent(child: right), newParent: findParent(child: left))
}

changeParent() 함수를 호출함과 동시에 양 쪽 인수마저 findParent() 함수로 각각 값을 확인하고 있다.

// 부모 바꾸기
func changeParent(oldParent:Int,newParent:Int) {
    //oldParent의 값을 가진 index를 찾아내서 해당 index의 값을 newParent로 바꿔준다.
    parents.indices.filter {parents[$0] == oldParent}.forEach{parents[$0] = newParent}
}

이 함수가 실행되려면 oldParent, newParent 값이 있어야 실행 될 수 있으니까?

//가장 위 부모 찾기
func findParent(child:Int) -> Int {
    //parrents의 child번째가 child라면 가장 위 부모 아니라면 재귀해줌
    return parents[child] == child ? child : findParent(child:parents[child])
}

이걸 두번 부르는데 이 함수는 하필 또 재귀다...!

즉, 과정이 너무 많다..!

 

그래서 이걸 줄일 수 있는 방법이 뭐가 있을지 생각해보다가

find 함수를 isCycle 다음에 바로 호출하고, 양 쪽 값 중 대소를 비교해서 change (union) 시켜주면 되지 않을까?

라는 생각이 들었고 실행에 옮겨 보았다.

 

import Foundation

var parents = [Int]()

func solution(_ n:Int, _ costs:[[Int]]) -> Int {
    
    parents = (0 ..< n).map{ $0 } // 각자 index의 값을 초기 값으로 설정
    let sortByCost =  costs.sorted { $0[2] < $1[2] } // cost 기준 오름차순 정렬
    var sum = 0 // 전체 합

    for bridge in sortByCost {
        if !isCycle(left: bridge[0], right: bridge[1]) {
            let first = find(child: bridge[0])
            let second = find(child: bridge[1])
            let minimum = min(first, second)
            let maximum = max(first, second)
            
            connect(oldParent:maximum,newParent:minimum)
            sum += bridge[2]
        }
    }
    return sum
}
// 사이클 형성 하는지
func isCycle(left:Int,right:Int) -> Bool {
    return parents[left] == parents[right] ? true : false
}

// 부모 바꾸기 (Union)
func connect(oldParent:Int,newParent:Int) {
    parents.indices.filter {parents[$0] == oldParent}.forEach{parents[$0] = newParent}
}

//가장 위 부모 찾기 (Find)
func find(child:Int) -> Int {
    return parents[child] == child ? child : find(child:parents[child])
}

 

그렇게 확인해본 결과!

왼쪽이 기존 코드, 오른쪽이 개선된 코드

확실히 시간 복잡도가 줄어든 걸 볼 수 있다!

특히 테스트 7에서 0.91ms 였는데 개선한 코드는 0.35ms로 많이 줄어들었다 ㅎㅎㅎ!

 

이제 Union-Find 알고리즘은 절대 안까먹을 것 같다ㅋㅋㅋㅋㅋㅋㅋ

결과를 보니 재밌기도 하고 뿌듯하기도 하다. 앞으로 코드 개선 경험을 늘려야겠다.

 

그래서 과제를 제출했지만?

더 짧은 방법이 있었다...

import Foundation

var parents = [Int]()

func solution(_ n:Int, _ costs:[[Int]]) -> Int {
    
    parents = (0 ..< n).map{ $0 } // 각자 index의 값을 초기 값으로 설정
    let sortByCost =  costs.sorted { $0[2] < $1[2] } // cost 기준 오름차순 정렬
    var sum = 0 // 전체 합

    for bridge in sortByCost {
        if union(bridge[0], bridge[1]) {
            sum += bridge[2]
        }
    }
    return sum
}

// 부모 바꾸기 (Union)
func union(_ a: Int, _ b: Int) -> Bool {
    let a = find(a)
    let b = find(b)
    if a == b {
        return false
    }
    if a > b {
        parents[a] = b
    } else {
        parents[b] = a
    }
    return true
}

//가장 위 부모 찾기 (Find)
func find(_ i:Int) -> Int {
    guard parents[i] != i else { return i }
        parents[i] = find(parents[i])
    return parents[i]
}