22996번: 유니온 파인드 복원

문제 설명

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},⋯,{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≤N≤300000
•  1≤M≤Q≤300000
• 1번 질의에서 1≤u,v≤N
• 1번 질의에서 u,v는 같을 수 있다.
• 2번 질의에서 1≤v≤N
• 준원이가 기억하는 정보와 일치하는 입력 데이터가 존재한다.

[예제 입력 1]

6 7
4 4 4 6 6 6
2
4
6

[예제 출력 1]

6 7
1 1 2
1 2 3
1 3 4
2 2
1 5 6
1 1 6
2 6

---

풀이

 풀이를 최대한 단순하게 만들려면, 결과를 보고 자연스러운 과정을 예측하려 하기 보다는 단순 결과에 맞게 답을 우겨 넣어야 한다.

 우선 가지고 있는 정보는 최종 부모 노드를 가리키는 배열과 2번 질의에 대한 답인데, 2번 질의는 매우 쉽게 해결이 가능하다. 가장 초기 해당 배열은 자신을 가리키고 있는 배열이고, 이 상태에서는 어떤 2번 질의든 해결할 수 있다. 예를 들어 2번 질의 답이 4인 경우 단순히 초기 배열에서 2 4를 하면 되기 때문이다. 이렇게 주어진 질의 중 2번 질의는 초반에 모두 출력을 해버리면 된다.

 이후 1번 질의에 대한 해결은 어렵다. 단순히 결과에 맞게 1번 질의를 우겨 넣어야 하는데, 그런 경우에는 최종 부모 배열을 가지고 바로 1번 질의를 실행해야 한다. 예를 들어 예제에서는 1 1 2, 1 2 3, 1 3 4 이후 1 1 6 질의를 통해 1번, 2번, 3번 노드의 부모로 4가 되는데, 단순하게 출력에서는 1 1 4, 1 2 4, 1 3 4라고 가정해야 한다. 하지만 이런 해결을 하는데도 고려할 부분이 있는데, 바로 상위 노드의 부모 노드가 없어야 한다는 점이다. 만약 1 1 4, 1 2 4, 1 3 4 질의를 하는데 이미 4번 노드의 부모가 6으로 되어 있으면 1, 2, 3번 노드 부모 모두 6이 되어 틀리게 된다. 그러기 때문에 어떤 노드 a에 대해 a를 다른 노드와 merge하기 전에 a를 부모 노드로 잡고 있는 서브 트리가 이미 완성되어 있어야 한다. 이는 최종 배열을 가지고 일종의 자식 트리를 가리키는 배열을 만들어, 최종 배열 노드들 중 루트 노드에서 dfs 탐색을 돌려 이를 배열에 넣어 놓아야 한다.

 그리고 이렇게 1번 질의를 직접적으로 하는 개수는 입력을 통해 알 수 있는 1번 질의의 개수보다 항상 작거나 같기 때문에, 이러한 추가적인 개수는 단순하게 1 1 1 질의를 출력함으로써 해결할 수 있다.

#include <iostream>
#include <vector>

using namespace std;

vector<int> child[300001];//자식 노드 저장
vector<int> rootNodes;//최상단 노드 저장
vector<pair<int, int>> query1;//쿼리 1 저장
vector<int> query2Result;//쿼리 2 결과 저장

void dfs(int num) {//dfs 기반 트리 탐색
	for (int i = 0; i < child[num].size(); i++) {
		query1.push_back({ child[num][i],num });

		if (child[child[num][i]].size() > 0)
			dfs(child[num][i]);
	}
}

int main() {
	int N, Q, M, temp;
	int changeCount = 0;

	cin >> N >> Q;

	for (int i = 1; i <= N; i++) {//결과 배열에서 최상단과 자식 노드 따로 저장
		cin >> temp;
		if (i == temp)
			rootNodes.push_back(temp);
		else {
			child[temp].push_back(i);
			++changeCount;
		}
	}

	cin >> M;

	for (int i = 0; i < M; i++) {
		cin >> temp;
		query2Result.push_back(temp);
	}

	for (int i = 0; i < rootNodes.size(); i++) //최상단 노드부터 탐색
		dfs(rootNodes[i]);

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

	for (int i = 0; i < M; i++)
		cout << "2 " << query2Result[i] << "\n";

	for (int i = 1; i <= Q - M - changeCount; i++)
		cout << "1 1 1\n";

	for (int i = query1.size() - 1; i >= 0; i--)
		cout << "1 " << query1[i].first << " " << query1[i].second << "\n";

	return 0;
}

---

결과