취미로 코딩하기

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

이번에 볼 문제는 백준 24452번 문제인 交易計画 (Trade Plan)이다.
문제는 아래 링크를 확인하자.

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

그래프가 주어질 때, 두 노드가 같은 component에 들어있는지 확인하는 작업은 disjoint set 자료구조를 이용해 간단히 할 수 있다. 또한 disjoint set을 rank compression을 기반으로 구현하면 이전 단계에 union한 두 component를 합치기 전으로 빠르게 돌려놓을 수 있게 된다. 이를 이용해 문제를 해결해보자.

구체적인 방법은 다음과 같다. 초기상태로 한 나라의 도시끼리를 잇는 도로들을 미리 전처리해두자. 그리고 나라A의 도시 x와 나라B의 도시 y가 나라 A와 B의 도시만을 거쳐 연결되어있는지 묻는 쿼리들을 몰아 처리해보자. 즉, 나라A와 나라B의 도시를 잇는 에지정보를 이용하여 그래프의 component들을 union하고 쿼리를 처리하자. 쿼리의 처리가 끝났다면 위 과정에서 변형된 그래프를 초기상태(한 나라의 도시끼리 잇는 도로들만 연결된 상태)로 다시 다시 리셋시키자.

쿼리의 순서를 조정하면, 각 에지가 그래프의 component를 합치고 다시 리셋하는 데에 사용되는 횟수는 많아야 1번이 되게 할 수 있다는 점을 관찰하자.

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

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

int uf[400001];
int depth[400001];
vector<pair<pair<int, int>, int>> stk;

int findf(int cur) {
	if (cur == uf[cur]) return cur;
	return findf(uf[cur]);
}

void unionf(int x, int y) {
	x = findf(x), y = findf(y);
	if (x == y) return;
	if (depth[x] < depth[y]) {
		uf[x] = y;
		stk.emplace_back(make_pair(make_pair(x, y), 0));
	}
	else if (depth[x] > depth[y]) {
		uf[y] = x;
		stk.emplace_back(make_pair(make_pair(y, x), 0));
	}
	else {
		uf[x] = y;
		depth[y]++;
		stk.emplace_back(make_pair(make_pair(x, y), 1));
	}
}
void unionf(int x, int y, int& cnt) {
	x = findf(x), y = findf(y);
	if (x == y) return;
	cnt++;
	if (depth[x] < depth[y]) {
		uf[x] = y;
		stk.emplace_back(make_pair(make_pair(x, y), 0));
	}
	else if (depth[x] > depth[y]) {
		uf[y] = x;
		stk.emplace_back(make_pair(make_pair(y, x), 0));
	}
	else {
		uf[x] = y;
		depth[y]++;
		stk.emplace_back(make_pair(make_pair(x, y), 1));
	}
}

int color[400001];

int N, M, K, Q;
pair<int, int> edges[400000];
vector<pair<int, int>> evec;
int ans[400000];
vector<pair<pair<int, int>, int>> qvec;

bool ecomp(pair<int, int>& p1, pair<int, int>& p2) {
	return color[p1.first] < color[p2.first] || (color[p1.first] == color[p2.first] && color[p1.second] < color[p2.second]);
}
bool eequal(pair<int, int>& p1, pair<int, int>& p2) {
	return (color[p1.first] == color[p2.first]) && (color[p1.second] == color[p2.second]);
}
bool qcomp(pair<pair<int, int>, int>& p1, pair<pair<int, int>, int>& p2) {
	return ecomp(p1.first, p2.first);
}

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

	cin >> N >> M >> K;
	for (int i = 1; i <= N; i++) uf[i] = i;

	for (int i = 0; i < M; i++) cin >> edges[i].first >> edges[i].second;
	for (int i = 1; i <= N; i++) cin >> color[i];
	for (int i = 0; i < M; i++) {
		if (color[edges[i].first] == color[edges[i].second]) unionf(edges[i].first, edges[i].second);
		else {
			if (color[edges[i].first] > color[edges[i].second]) swap(edges[i].first, edges[i].second);
			evec.emplace_back(edges[i]);
		}
	}
	sort(evec.begin(), evec.end(), ecomp);

	cin >> Q;
	for (int i = 0; i < Q; i++) {
		int x, y; cin >> x >> y;
		if (color[x] == color[y]) ans[i] = (findf(x) == findf(y)) ? 1 : 0;
		else {
			if (color[x] > color[y]) swap(x, y);
			qvec.emplace_back(make_pair(make_pair(x, y), i));
		}
	}
	sort(qvec.begin(), qvec.end(), qcomp);

	int esize = evec.size(), qsize = qvec.size();
	int eidx = 0, qidx = 0;
	while (eidx < esize && qidx < qsize) {
		auto& e = evec[eidx]; auto& q = qvec[qidx];
		if (ecomp(e, q.first)) eidx++;
		else if (ecomp(q.first, e)) qidx++;
		else {
			int cnt = 0;
			while (eidx < esize && eequal(e, evec[eidx])) {
				unionf(evec[eidx].first, evec[eidx].second, cnt);
				eidx++;
			}
			while (qidx < qsize && eequal(q.first, qvec[qidx].first)) {
				ans[qvec[qidx].second] = (findf(qvec[qidx].first.first) == findf(qvec[qidx].first.second)) ? 1 : 0;
				qidx++;
			}
			while (cnt--) {
				auto& lglg = stk.back();
				uf[lglg.first.first] = lglg.first.first;
				if (lglg.second) depth[lglg.first.second]--;
				stk.pop_back();
			}
		}
	}

	for (int i = 0; i < Q; i++) cout << ans[i] << '\n';
}