[백준 C++]1717번: 집합의 표현(서로소 집합. union & find)

1.문제

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

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

2. Solved.ac 난이도

3. 활용 알고리즘:  Union & Find(그래프)

4. 틀린 이유: TLE

진짜 왠만하면 endl대신 \n을 써야한다는 것을 뼈져리게 느꼈다.
scanf는 불편해서 일단 보류한다고 쳐도, endl과 \n이 그렇게 시간차이가 큰지 몰랐음.

---

시도1. (TLE)

처음엔 그래프형 union& find를 써야한다는 생각을 하지 못했다. 

때문에 일단 brute force로 아무렇게나 구현해봤는데 역시나 시간초과..

#include<iostream>
#include <vector>
#include<algorithm>
using namespace std;
vector<vector<int>> set;
int have(int a) {
	for (int i = 0; i < set.size(); i++) {
		if (find(set[i].begin(), set[i].end(), a) != set[i].end()) {
			return i;
		}
	}
	return -1;
}
void uni(int a, int b) {
	int v1 = have(a);
	int v2 = have(b);
	if (v1 == v2 && v1 != -1) {
		return;
	}
	else if (v1 == -1 && v2 == -1) {
		set.push_back({ a,b });
	}
	else if (v1 == -1) {
		set[v2].push_back(a);
	}
	else if (v2 == -1) {
		set[v1].push_back(b);
	}
}

int main() {
	
	int n, m,op,val1,val2;
	cin >> n >>  m;
	for (int i = 0; i < m; i++) {
		cin >> op>>val1>>val2;
		if (op == 1) {
			int v1=have(val1);
			int v2 = have(val2);
			if (v1 == v2&&v1!=-1) {
				cout << "YES";
			}
			else {
				cout << "NO";
			}
		}
		else {
			uni(val1, val2);
		}
	}
	return 0;
}

 

시도2 (TLE)

다시 구현해보았는데도 시간초과가 났다. 이유는 find함수. 처음에 find함수를 저렇게 재귀로 구현해놓으니까 너무 시간낭비가 심했다.

ex) par[3]->par[2]->par[8]->par[13]->par[39]=39 가 될 수도 있는 것이다. 

#find 함수부분
int find(int a) {
	if (parent[a] == 0) {
		parent[a] = a;
	}
	if (parent[a] == a)
		return a;
	return find(parent[a]);
}

 

 

최종 정답 코드

따라서 아래와 같이 바꿧다!

#include<iostream>
#include <vector>
using namespace std;
int parent[1000001];

int find(int a) {
	if (parent[a] == a)
		return a;
	else {
		parent[a] = find(parent[a]);
	}
	return parent[a];
}
void uni(int a, int b) {
	int pa = find(a);
	int pb = find(b);
	parent[pa] = pb;
}

int main() {
	ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
	int n, m, op, val1, val2;
	cin >> n >>  m;
	for (int i = 0; i <= n; i++) {
		parent[i] = i;
	}
	for (int i = 0; i < m; i++) {
		cin >> op>>val1>>val2;
		if (op == 1) {
			if (find(val1)==find(val2)) {
				cout << "YES\n";
			}
			else {
				cout << "NO\n";
			}
		}
		else {
			uni(val1, val2);
		}
	}
	return 0;
}