Disjoint-Set / Union-Find (디스조인트 셋, 유니온 파인드)

저는 예전에 디스조인트 셋이 유니온 파인드인지 몰라서 엥? 디스조인트 셋? 이건 뭐랭 했었는데 알고보니 같은 말이었네요ㅋㅋㅋ 그래서 말이죠, 오늘은 크루스칼에서도 사용되는 Disjoint Set을 알아볼게요. 

 

Disjoint Set의 원리는 서로소 집합과 연관되어 생각하면 됩니다. 새로운 집합을 추가하거나, 기존의 집합과 합치거나, 집합에 어떤 원소가 들어있는지 찾는 연산을 거의 선형시간에 가깝게 수행하는 자료구조라고 할 수 있죠.

Disjoint Set이 Union-Find라고도 불리는 이유는, Disjoint Set을 구성하는데에 union과 find라는 함수? 개념? 을 사용하기 때문입니다. find를 통해서 a라는 요소가 b라는 요소와 같은 집합인지(연결되어 있는지) 를 확인 할 수 있고, union을 통해 두 요소를 연결시켜서 집합을 합치도록 할 수 있습니다.

 

즉, Disjoint Set에서 부분집합들은 tree로 표현될 수 있고, find를 통해 루트 노드를 찾아 union으로 tree를  합치도록 합니다.

• find : 어떤 원소가 주어졌을 때, 이 원소가 속한 집합을 반환합니다. 정확히 말하면, 어떤 원소가 속한 집합을 대표하는 원소를 반환합니다.
• union : 두 개의 집합을 하나의 집합으로 합칩니다.

 

그럼 가장 기본적인 Disjoint Set을 swift로 구현해보겠습니다.

 

[1] 먼저, 자기 자신을 root로 두어 초기화합니다. (처음에는 자기 자신이 루트 노드가 되는 거겠죠?)

 

[2] 주어진 원소의 루트 노드 번호를 반환하는 find를 구현합니다. (자신이 속한 집합(트리)에서 대표하는 원소를 반환합니다!)

 

[3] 두 원소를 합치는 union을 구현합니다. 먼저 각 원소의 루트 노드를 찾고, 하나의 루트 노드를 다른 하나의 루트 노드로 바꿉니다.

 

---

 

생각보다 코드가 간결하고 쉽죠? 하지만 응용단계로 넘어가면 조금씩 헷갈리기 시작합니다ㅠㅠ 막 union 과정에서 뭘 계산해줘야하고, 지금 제가 설명한 union과 다르게 aRoot와 bRoot가 아닌 a와 b를 연결하도록 하기도 하죠. 그치만 이건 문제에 따라 어떻게 푸는지 프로세스를 잘 생각해보고 구현한다면 문제가 되지 않을 거라고 생각합니다 :)

 

그래서 마지막으로 예제를 하나 보려고 합니다. 예제는 백준 3780 - 네트워크 연결 문제입니다.

 

 

예제 입력을 보면 최종적으로는 3-1-2-4 로 기업이 이어지게 됩니다. 문제 설명을 보면 A-B로 이어지는 클러스터는 B의 센터에 의해 제공된다고 했기 때문에, E 3을 입력하면 4로부터의 거리인 5가 출력되게 되죠.

 

그럼 위 문제를 읽어봤을 때 알 수 있는 점은, 입력된 문자의 첫 글자가 E인 경우 거리를 출력해주고, I인 경우 두 기업을 새로운 클러스터로 이어주어야 합니다.

 

위에서 Disjoint Set의 기본 구현을 했을 때는 각 요소들의 루트 노드를 가리키는 parent 배열만 필요로 했었는데, 이번 문제는 각 기업과 센터 간의 거리를 나타내는 distance 배열도 필요로 하겠죠?

 

그리고 자신의 루트 노드를 find할 때 자신의 distance에 부모의 distance를 더하도록 하여 센터(루트 노드)까지의 거리를 구하도록 합니다.

 

여기서 만약 4의 부모가 2, 2의 부모가 1(센터)이라면 이 find 함수는 어떻게 동작할까요? 👉 find(4)

tmp = find(parent[value])로 재귀가 발생하기 때문에 아래 distance 어쩌구 코드는 4부터가 아닌 2부터 실행됩니다.

 

처음에는

distance[2] += distance[1] (distance[1] 은 여기서 센터기 때문에 0입니다.)

distance[4] += distance[2]

를 통해 2에서 1까지의 거리, 4에서 1까지의 거리를 구할 수 있게 됩니다.그리고 이 때, 4의 부모도 1로 바꿔주게 되죠.

 

그럼 이후 find(4) 함수를 다시 실행시키더라도, 4의 부모는 1이 되었기 때문에 distance[4] += distance[1] (distance[1] 은 여기서 센터기 때문에 0입니다.) 이 되어 아무런 영향을 끼치지 않고 정상적으로 distance[4]를 구해주게 됩니다.

 

따라서 E 라는 키워드로 거리를 출력해줄 때, find를 사용하여 혹시 계산되지 않은 거리를 해주도록 하고 출력하면 됩니다.

 

그럼 I 라는 키워드로 두 기업을 연결시키려면 union 함수가 있어야 겠죠? union 함수는 좀 더 간단합니다. 한 쪽의 부모를 다른 쪽으로 지정하여 연결하도록 합니다. 그리고 a와 b의 거리를 더해주면 되죠.

 

 

 

어느정도 이해가 가셨나요? 혹시 아직 잘 모르신다고 해도 아래 전체 코드를 통해 중간중간 print로 과정을 출력해서 보다보면 어느정도 이해하실 수 있을 것 같습니다😄

 

전체 코드는 아래에 있습니다 :)

let test_case = Int(readLine()!)!

var parent: [Int] = []
var distance: [Int] = []
for _ in 1...test_case {
    let company = Int(readLine()!)!
    parent = Array(repeating: 0, count: company+1)
    distance = Array(repeating: 0, count: company+1)
    for i in 1...company { parent[i] = i }
    
    var input = readLine()!
    while input != "O" {
        let order = input.split(separator: " ").map { String($0) }
        if order[0] == "E" { printDistance(Int(order[1])!) }
        else if order[0] == "I" { union(Int(order[1])!, Int(order[2])!) }
        input = readLine()!
    }
}

func printDistance(_ value: Int) {
    find(value)
    print(distance[value])
}

func union(_ a: Int,_ b: Int) {
    parent[a] = b
    distance[a] += abs(a - b) % 1000
}

func find(_ value: Int) -> Int {
    if parent[value] == value { return value }
    
    let tmp = find(parent[value]);
    distance[value] += distance[parent[value]];
    parent[value] = tmp
    
    return parent[value]
}