골드4 : 유니온 파인드 문제이다.

생각

문제를 이해해보자. 먼저 처음에 아무일도 하지 않을 경우에 입력된 n에 대해서 0 ~ n까지의 바구니가 생긴다. 그 다음, m개의 명령이 들어온다. 명령은 해당 원소가 들어있는 바구니를 합치는 것, 그리고 해당 원소가 같은 바구니에 들어있는 지 확인하는 것이다.

그렇다면 핵심은, 현재 어떻게 내가 들어가 있는 바구니를 찾을 것인가? 그리고 찾았다면 어떤식으로 합쳐서 가지고 있을 것인가? 이다. 결국 집합을 만들고, 서로소 인지를 확인하는 문제이다.

가장 단순하게 생각해보자. 처음에 모든 바구니는 자기자신만을 원소로 갖는 바구니를 가진다. 이 상태에서 1번 바구니와 3번 바구니를 합치라는 명령이 내려올 때, 우리는 일반적으로 **이것을 동등하게 생각하고 연결**하려 한다. 하지만 이렇게 생각할 경우에 해당 바구니에 대한 번호를 재정의 해야 한다. 예를 들어, (1, 2)가 들어있는 바구니는 몇번 바구니라 정의해야 할까? 또 추가적으로 3이 들어온 경우 (1, 2, 3)은 몇번 바구니 일까? 바구니에 원소가 추가될 경우 계속하여 바구니의 이름을 재정의 해야 한다.

그러니 위계 질서를 줘보자. 1번 바구니와 2번 바구니를 합친다는 명령이 내려왔을 때 2번 바구니는 1번을 따른다. (21) 그리고 3번 바구니는 2번을 따른다고 하면 (321)로 준다. 그리고 이렇게 생성된 바구니의 번호를 1이라 하면, 합치는 과정을 통한다하더라도 합쳐진 집합의 번호는 가장 먼저 합쳐진, 즉 조상 원소의 번호라 지칭할 수 있다. 즉, 위의 바구니는 1번 바구니라 할 수 있고, 다른 합쳐진 집합이 있다하더라도 이 집합의 조상은 독립적으로 유지된다.

(3->2->1) = 1번 바구니
(4->5) = 5번 바구니

5번 바구니와 3번 바구니를 합해라.
=> (4->5->3->2->1) = 1번 바구니

문제 발생

하지만 이 경우에는 문제가 발생하는데, (21), (54) 인 경우가 있다고 하자. 이 때 5번이 2번 바구니에 속한다고 할 경우 이 논리대로라면 (521)이 된다. 그런데 사실은 4번도 1번 바구니에 속해야 한다. 논리적 오류가 발생한다.

이 문제는 바구니가 속하는 기준이 한 방향이기 때문이다. 하지만 속하는 기준을 여러개로 만들지 않아도 해결이 가능하다.

해결 방법

아까 조상을 만들고, 추가된 집합에 대한 바구니 번호를 가장 위에 있는 조상의 번호를 따르기로 했다. 위에서 발생한 문제에 대해 해당 번호가 어떤 바구니(조상)에 있는지 확인하고 그 바구니끼리 연결한다면 해당 문제는 해결된다.

(2->1) = 1번 바구니
(5->4) = 4번 바구니

5번 바구니와 2번 바구니를 합쳐라.

5번의 조상 : 4번
2번의 조상 : 1번

(4->1)로 바꿔버린다.

=> (2->1), (5->4->1)

이렇게 된 경우 1번 바구니에 해당한 번호는 2, 4, 5로 원하는 결론이다.

구현

이것을 구현하는 것이 어려워 보이지만, 단순하다. 위에서 결국 필요한 것은 내 원소의 번호, 그리고 그 원소가 따르는 번호이다. 그리고 그 따르는 것들의 바구니 번호는 가장 위에 기초가 되는 조상 번호가 대표한다.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|-| |번호|0|1|2|3|4|5|6|7|

7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1

line 1

0 1 3 // 1번과 3번을 합쳐라.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|3|4|5|6|7|

1의 조상 번호 : 1
3의 조상 번호 : 3
3번이 따르는 바구니를 1로 바꾼다.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|6|7|

line 2

1 1 7 // 1번과 7번은 같은 바구니이니?

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|6|7|

7의 조상 번호 : 7
1의 조상 번호 : 1

No

line 3

0 7 6 // 7번과 6번을 합쳐라.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|6|7|

7의 조상 번호 : 7
6의 조상 번호 : 6
6번은 7번을 따른다.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|7|7|

line 4

1 7 1 // 7번과 1번은 같은 바구니이니?

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|7|7|

7의 조상 번호 : 7
1의 조상 번호 : 1

No

line 5

0 3 7 // 3번과 7번을 합쳐라.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|7|7|

3의 조상 번호 : 1
7의 조상 번호 : 7
7번은 3번을 따른다.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|7|3|

line 6

0 4 2 // 4번과 2번을 합쳐라.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|7|3|

4의 조상 번호 : 4
2의 조상 번호 : 2
2번은 4번을 따른다.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|4|1|4|5|7|3|

line 7

0 1 1 // 1번과 1번을 합쳐라.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|2|1|4|5|7|3|

1의 조상 번호 : 1
1의 조상 번호 : 1
1번은 1번을 따른다.

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|4|1|4|5|7|3|

line 8

1 1 1 // 1번과 1번은 같은 바구니이니?

|바구니|0|1|2|3|4|5|6|7| |:-:------|:|:|:|:|:|:|:|:| |번호|0|1|4|1|4|5|7|3|

1의 조상 번호 : 1
1의 조상 번호 : 1

Yes

결론

이런 방식으로 업데이트 하는 과정을 거치면, 묶인 집합을 효과적으로 찾을 수 있다. 결론적으로 구현해야 하는 함수는 find, union 함수이다.

find

find 함수같은 경우에는, 특정 원소에 대해 물었을 때, 그 녀석이 속한 바구니를 끝까지 찾아 어떤 바구니인지 반환한다.

union

union 함수는 명령이 들어왔을 때, 각각이 속해있는 바구니를 찾고 그 바구니를 연결한다. 이 때, 중요한 것은 해당 원소의 바구니 번호를 찾아야 한다는 것이다.

#include <iostream>
#include <algorithm>
#include <string>
#include <vector>
using namespace std;
int n, m;
int box[1000001];
 
int find(int elem){
    if (elem == box[elem]) {
        return elem;
    }
    return find(box[elem]);
}
void merge(int elem1, int elem2){
    box[find(elem2)] = find(elem1);
}
 
int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
 
    cin >> n >> m;
    for (int i = 0; i <= n; i++) {
        box[i] = i;
    }
    for (int i = 0; i < m; i++) {
        int command, elem1, elem2;
        cin >> command >> elem1 >> elem2;
        if (command) {
            if (find(elem1) == find(elem2)) {
                cout << "YES" << '\n';
            } else {
                cout << "NO" << '\n';
            }
        } else {
            merge(elem1, elem2);
        }
    }
    return 0;
}

시간 복잡도 문제

하지만 위의 코드대로 한다면 문제가 발생한다. 그 이유는 원소의 개수가 많아짐에 따라 어떤 바구니에 들어있는지 탐색하는 시간이 오래걸리기 때문이다. 예제 입력에 대한 마지막 상황을 보자.

|바구니번호|0|1|2|3|4|5|6|7| |:---:----------|:|:|:|:|:|:|:|:| |따르는 번호|0|1|4|1|4|5|7|3|

이 상황에서, 6은 어느 바구니에 있을까? (6731) 로, 1번 바구니에 있다. 만약, 6번을 0번이 속한 바구니와 합치라는 명령이 들어온다면 4번의 연산을 수행해야 한다. 그렇다면, 극단적인 경우에, n개의 원소가 순차적으로 얽혀있고 (nn-1n-2…) n이라는 원소를 특정 바구니에 합치라는 명령이 들어올 경우, 의 시간 복잡도를 갖는다. 이 때, 수행하는 명령의 개수가 m일 경우 의 시간 복잡도를 갖는다. 현재 문제에서 n = 100000, m = 1000000이므로 시간 초과가 난다.

해결

이 문제는 해결이 간단한데, find 함수에 찾는 과정에서, 최종 조상의 번호를 아예 해당 원소의 바구니라고 저장해버리는 것이다. 이 상황에서는 원소가 속해있는 바구니 번호가 중요한 것이지, 그 과정에서 발생하는 그래프가 중요한 것이 아니기 때문이다.

#include <iostream>
#include <algorithm>
#include <string>
#include <vector>
using namespace std;
int n, m;
int box[1000001];
 
int find(int elem){
    if (elem == box[elem]) {
        return elem;
    }
    return box[elem] = find(box[elem]);
}
void merge(int elem1, int elem2){
    box[find(elem2)] = find(elem1);
}
 
int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
 
    cin >> n >> m;
    for (int i = 0; i <= n; i++) {
        box[i] = i;
    }
    for (int i = 0; i < m; i++) {
        int command, elem1, elem2;
        cin >> command >> elem1 >> elem2;
        if (command) {
            if (find(elem1) == find(elem2)) {
                cout << "YES" << '\n';
            } else {
                cout << "NO" << '\n';
            }
        } else {
            merge(elem1, elem2);
        }
    }
    return 0;
}

Reference