112. 유니온 파인드 복원 (백준 22996) [C++]

 문제

문제 링크 : https://www.acmicpc.net/problem/22996

22996번: 유니온 파인드 복원

문제

준원이는 스타팅 포인트의 멤버이다. 준원이는 회사에서 병합 (merge) 연산을 수행할 때마다 포인트를 얻는다.

그래서 준원이는 C언어로 다음과 같은 코드를 작성했다.

#include <stdio.h>

int par[300001];
int find(int v) {
    if (v == par[v]) return v;
    return par[v] = find(par[v]);
}
void merge(int u, int v){
    u = find(u);
    v = find(v);
    if (u != v) par[u] = v;
}

int main() {
    int N, Q;
    scanf("%d %d", &N, &Q);
    for (int i=1; i<=N; i++) par[i] = i;

    for (int i=1; i<=Q; i++) {
        int type;
        scanf("%d", &type);
        if(type == 1) {
            int u, v;
            scanf("%d %d", &u, &v);
            merge(u, v);
        }
        else if(type == 2) {
            int v;
            scanf("%d", &v);
            printf("%d\n", find(v));
        }
    }
}

준원이는 위 코드로 다음 문제를 풀었다.

 $1$부터 $N$까지 $N$개의 자연수를 원소로 가지는 $N$개의 집합 ${1}, {2}, \cdots , {N}$이 있다.

여러분은 아래와 같이 주어지는 질의를 순서대로 $Q$개 처리해야 한다.

• 1번 질의: 1 u v 의 형식으로 주어진다. 원소 $u$가 속한 집합과 원소 $v$가 속한 집합이 다르다면, 두 집합을 하나로 합친다.
• 2번 질의: 2 v 의 형식으로 주어진다. 원소 $v$가 속한 집합에 있는 원소를 아무거나 하나 반환한다.

그런데, 문제가 생겼다. 준원이는 어떤 입력 데이터를 받았는지 까먹었다. 하지만 주어진 $N, Q$ 값과, 2번 질의들이 반환한 값들, 그리고 실행 결과 최종 상태에서의 par[] 배열의 값들은 기억하고 있다.

이 정보를 이용해 어떤 입력 데이터가 주어졌을지 복원해보자.

입력

첫째 줄에 준원이가 기억하는 원소의 개수 $N$과 질의의 개수 $Q$가 공백으로 구분되어 주어진다.

둘째 줄에 실행 결과 최종 상태에서의 par[] 배열의 값들을 의미하는 $N$개의 자연수가 공백으로 구분되어 주어진다.

셋째 줄에 입력 데이터에서 주어진 2번 질의의 개수 $M$이 주어진다.

넷째 줄부터 $M$개의 줄에 걸쳐, 준원이의 코드가 2번 질의들이 반환한 값들이 한 줄에 하나씩 순서대로 주어진다.

출력

준원이의 프로그램이 받았을 입력 데이터를 아래와 같은 형식으로 복원하여 출력한다.

첫째 줄에 원소의 개수 $N$과 질의의 개수 $Q$를 공백으로 구분하여 출력한다.

둘째 줄부터 $Q$개의 줄에 걸쳐, 질의를 한 줄에 하나씩 순서대로 출력한다.

• 1번 질의: 1 u v 의 형식으로 출력한다.
• 2번 질의: 2 v 의 형식으로 출력한다.

준원이가 기억하는 정보와 일치하는 입력 데이터가 여러 가지 있으면 아무거나 출력해도 된다.

제한

•  $1 \leq N \leq 300,000$ 
•  $1 \leq M \leq Q \leq 300,000$ 
• 1번 질의에서 $1 \leq u, v \leq N$ 
• 1번 질의에서 $u, v$는 같을 수 있다.
• 2번 질의에서 $1 \leq v \leq N$ 
• 준원이가 기억하는 정보와 일치하는 입력 데이터가 존재한다.

 

 풀이

JAVA로 시도하다 추가시간이 없어서 C++로 선회했다.

문제를 풀때 그냥 가장 간단한 방법들만을 찾으려고 했다.

 

우선 가장 생각하기 쉬운 2 v형식을 살펴보자.

최초의 상태에선 모든 i는 각각의 배열에 들어있으므로 -> 그냥 자기 자신을 출력하면 된다.

2 v는 조합 형성에 영향을 주지도 않기 때문에 그냥 처음에 다 해결해버리면 된다.

 

그 다음으로 1 u v를 살펴봐야하는데, 문제의 예제입력과 코드를 잘 살펴봐야 한다.

예제를 확인하면, 4 4 4 6 6 6에서 결국 1,2,3 모두 6이 속한 배열인데 값은 4로 되어있다.

코드를 보면, find(v) 로 v의 끝까지 (루트까지) 찾고 -> if문을 통해 par[u]에 할당한다.

그러니깐, 1,2,3은 1 4 6 이전에 1 1 4 로 4에 할당되었고, 그 뒤에 할당되지 않았다는 것이다. (최 하단의 자식부터 먼저들 이동해야함)

 

그럼 이를 역으로 구현하기 위해서는

1. 배열의 값이 자기 자신과 같은 i를 찾아 추가하고 (루트)
2. 값이 i인 다른 j들부터 시작해서 j의 자식들을 j로 옮긴다.

이 순서를 거꾸로 출력하면 되기 때문에, stack을 사용했다.

 

그렇다면, 이렇게 1, 2번 규칙에서 가장 간단하만 해도 정답이 성립될까 (Q를 만족할까)

확신이 없기 때문에 문제의 제한을 잘 확인해야 한다.

1 u v에서 u와 v는 같을 수 있다. -> 1 1 1이라는 값이 존재할 수 있다는 뜻이다.

아무 변화가 없는 1 1 1을 집어넣어서 Q의 개수를 나중에 맞춰주면 되는것이다.

 

정리하면 출력의 순서는 다음처럼 하면 된다.

1. 2 v 질의를 무조건 본인으로 출력하게 해서 M개 맞추기 (출력)
2. Q - (1 u v를 수행했을때 결과 + M) 만큼 1 1 1 넣기
3. 1 u v 질의 출력

시간 초과를 많이 걸렸는데, C++에서 endl보다 '\n'이 훨신 빠르다는 것을 뒤늦게 알았다..

 

 코드

#include <iostream>
#include <vector>
#include <utility>
#include <stack>

using namespace std;

int N, Q, M;
vector<int> root;
stack<pair<int,int>> answer1;
vector<int> answer2;
vector<vector<int>> parSet(3000001, vector<int>(0,0));

void rollbackMerge(int a);

int main(int argc, char* argv[]) {

    cin >> N >> Q;

    int val;
    for(int i=1; i<N+1; i++){
        cin >> val;

        if(val == i){
            root.push_back(i);
        }else{
            parSet[val].push_back(i);
        }
    }

    cin >> M;
    //2번 값
    int a;
    for(int i=0; i<M; i++){
        cin >> a;
        answer2.push_back(a);
    }

    for(int i=0; i<root.size(); i++){
        rollbackMerge(root[i]);
    }

    cout << N << " " << Q << '\n';

    for(int s: answer2){
        cout << "2 " << s << '\n';
    }

    int lastSize = Q - answer1.size() - M;
    for(int i=0; i<lastSize; i++){
        cout<< "1 1 1" << '\n';
    }

    while(!answer1.empty()){
        cout<< "1 " << answer1.top().first << " "<< answer1.top().second << '\n';
        answer1.pop();
    }


    return 0;
}

void rollbackMerge(int a){
    for(int i: parSet[a]){

        answer1.push(make_pair(i,a));

        if(parSet[i].size() > 0){
            rollbackMerge(i);
        }
    }
}

 

 결과