취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 25200번 문제인 곰곰이와 자판기이다.
문제는 아래 링크를 확인하자.

https://www.acmicpc.net/problem/25200 

set 30만개의 배열을 준비하자. 각 i번째 set은 음료 i가 나오게 되는 초기 음료들의 집합이다.

쿼리 U, V가 주어질 때마다 U번째 set과 V번째 set을 합친 set을 V번째 set으로 바꿔주자. 이 때, 원소의 개수가 적은 set의 원소들을 원소의 개수가 많은 set에 추가하고 swap 등을 이용해 합쳐진 set이 V번째 set이 되게 하자. (물론 나머지 set은 clear()을 이용해 비워주자.)

위의 과정에서 각 원소가 다른 set으로 옮겨지는 경우, 해당 원소가 들어있는 set의 크기는 항상 두 배 이상이 되므로 각 원소는 많아봐야 O(lgN)번 옮겨담아지게 된다. 따라서 "원소를 한 set에서 다른 set으로 옮기는 횟수"는 많아야 O(NlgN)회 일어난다. 이는 문제 해결에 충분한 시간복잡도이다. (swap을 통해 두 집합을 바꾸는 작업은 swap의 시간복잡도가 O(1)이므로 위의 횟수를 셀 때 고려하지 않는다.)

아래는 제출한 소스코드이다.

#include <iostream>
#include <set>
#include <algorithm>
using namespace std;

set<int> st[300001];
int ans[300001];

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	int N, M; cin >> N >> M;
	for (int i = 1; i <= N; i++) {
		st[i].insert(i);
	}
	while (M--) {
		int u, v; cin >> u >> v;

		auto& stu = st[u], &stv = st[v];
		if (stu.size() <= stv.size()) {
			for (auto& n : stu) stv.insert(n);
			stu.clear();
		}
		else {
			for (auto& n : stv) stu.insert(n);
			stv.clear();
			swap(stu, stv);
		}
	}

	for (int i = 1; i <= N; i++) {
		for (auto& n : st[i]) ans[n] = i;
	}

	for (int i = 1; i <= N; i++) cout << ans[i] << ' ';
}

changelog

2022-05-16 내용 보충