백준 1717번 집합의 표현 - SWIFT

문제

---

초기에 {0}, {1}, {2}, ... {n} 이 각각 n+1개의 집합을 이루고 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

 

입력

첫째 줄에 n(1 ≤ n ≤ 1,000,000), m(1 ≤ m ≤ 100,000)이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다. a와 b는 n 이하의 자연수 또는 0이며 같을 수도 있다.

 

출력

1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다.

(yes/no 를 출력해도 된다)

 

문제 링크

1717번: 집합의 표현

 

 

풀이

---

처음에는 간단하게 생각해서 모든 배열을 만들어주는 방식으로 했다.

시간 복잡도고 뭐고 그냥 생각나는대로 만들어봤는데 예상대로 시간초과가 나왔다.

시간을 줄이기 위해서는 어떤 과정을 줄여야할까 생각을 하다가 아래 태그(분리 집합)를 통해 힌트를 얻었다.

나는 분리 집합의 존재를 몰랐기에 일단 분리 집합이 뭔지부터 알아봤다.

 

분리집합

교집합이 존재하지 않는 둘 이상의 집합을 말한다.

트리 구조는 부모 노드가 자식 노드를 가리키는 형태라면, 분리 집합은 자식 노드가 부모 노드를 가리키는 형태이다.

교집합이 없는 집합이기 때문에 차집합 연산은 의미가 없어 존재하지 않는다.

Union-Find 알고리즘을 사용한다.

 

합집합(Union) 연산

두 집합을 하나로 합친다.▼

 

 

이 때 우리가 현실에서 생각하는 것 처럼 배열에 집어넣는다거나 하는게 아니라,

C의 부모 노드를 A로 설정하는 것이다.

만약 두 원소의 최상위 부모 노드인 루트 노드가 같다면 둘은 같은 집합이다.

 

집합 탐색(Find)

해당 원소 노드가 속해 있는 집합을 찾는다.▼

 

 

루트 노드를 찾아서 둘을 비교하는 것이다.

더 자세한 내용은 아래의 링크에 정리되어있다.▼

분리집합과 Union-Find

 

 

 

프로그램 동작

분리집합 연산을 해주면 된다.

일단 노드 집합인 `parent[]` 배열을 만들었다.

초기에는 노드들은 각자 자기 자신을 가리키고 있다.

`parent[]` 배열의 값은 부모 노드를 의미한다.

`parent[]` 배열의 값이 현재 인덱스와 같다는 것은 자기 자신을 가리키고 있다는 것이고, 이는 최상위 노드라는 의미이다.

일단 명령으로 0이 들어오면, 합 연산을 시작한다.

합할 두 원소, `a`, `b`가 들어오면 각각의 부모 노드를 구하기 위해 집합 탐색을 한다.

부모 노드 탐색은 간단하다.

1. a값이 들어왔다고 하면, 현재 `parent[a]` 값을 구한다. 이 값이 곧 `a` 원소의 부모 노드 값이다.▼

 

 

2. `parent[a]`가 `a`와 같은지 비교해본다.

만약 같다면 이는 곧 a가 최상위 노드라는 뜻이므로 a를 반환한다.

만약 같지 않다면 parent[a] 값, 즉 a의 부모 노드 값을 가지고 노드 탐색 함수를 다시 돌린다. 1번으로 돌아가란 얘기다.▼

 

 

이 과정을 거쳐 a, b 각각의 루트 노드를 찾았다면 이제 두 루트 노드의 값을 비교한다.

일관성 없게 랜덤하게 두 루트 노드 중 아무거나 고르게 된다면 노드가 이리저리 꼬일 수 있어서 두 루트 노드 중 작은 값으로 합치게끔 한다.

만약 a 루트 노드가 1, b 루트 노드가 6이라면 1쪽으로 합친다는 것이다.

합치는건 간단하다.

b 루트 노드인 3의 루트 노드를 1로 만들어주면 된다.▼

 

 

이런 식으로 합쳐주면 실제로 배열에 넣고 검사하고 하는 과정을 안해도 된다.

 

 

Swift 코드

---

let nm = readLine()!.split(separator: " ").map { Int(String($0))! }
let n = nm[0], m = nm[1]
var parent = [Int]()

for i in 0...n {
    parent.append(i)
}

func findParent(_ a: Int) -> Int {
    if a == parent[a] {
        return a
    } else {
        parent[a] = findParent(parent[a])
        return parent[a]
    }
}

func makeUnion(_ a: Int, _ b: Int) {
    let pa = findParent(a);
    let pb = findParent(b);

    if pa > pb {
        parent[pa] = pb
    } else if pa < pb {
        parent[pb] = pa
    }
}

for _ in 0..<m {
    let command = readLine()!.split(separator: " ").map { Int(String($0))! }

    if command[0] == 0 {
        makeUnion(command[1], command[2])
    } else {
        if findParent(command[1]) == findParent(command[2]) {
            print("YES")
        } else {
            print("NO")
        }
    }
}

 

Swift에서는

return findParent(parent[a])

 

는 가능해도,

return parent[a] = findParent(parent[a])

 

같은 C문법은 사용할 수 없으니 주의해서 코딩해야한다.