[C++] 백준 25587번: 배수로

문제 링크

25587번: 배수로

문제 요약

N개의 도시에 대해서 각각 배수로 용량과 강수량이 주어진다. 강수량이 배수량보다 크다면 홍수가 발생하게 된다.
만약, 도시 사이에 배수로를 놓으면, 도시는 강수량과 배수량은 합쳐져서, 연결된 모든 도시에 홍수가 나거나, 나지 않게 된다.
이제, 다음 두 개의 쿼리가 M개 주어질 때, 해당 쿼리들을 수행하는 프로그램을 작성해야 한다.
1. 두 도시 사이에 배수로를 연결한다.
2. 현재 홍수가 날 도시의 수를 출력한다.

접근 방법

현재 기준으로 정답률이 85%를 넘는, 간단한 분리 집합 문제입니다.

도시를 연결하고, 강수량과 배수량을 계산하는 것은 분리 집합을 이용하면 간단하게 처리할 수 있습니다.

먼저 배수로 용량(a 배열)과 강수량(b 배열)을 입력 받으면서, '홍수가 발생할 수 있는 도시의 수'를 전역변수에 세줍니다. 분리 집합을 사용할 것이기 때문에 parent 배열을 만들어 두고, 모두 인덱스와 동일한 값으로 초기화해줘서 각각의 도시를 집합으로 만들어 줍니다.

추가적으로, 도시의 수와 동일한 크기를 가진 cnt 배열을 만들어 주고, 모두 1로 초기화 해줍니다. 이는 각 집합의 크기를 저장하는 용도로 사용할 예정입니다.

1번 쿼리는 다음과 같이 처리해주면 됩니다.

1. Find() 함수를 통해서 두 도시가 같은 집합에 속해있는지 확인해줍니다.
2. 같은 집합이라면 넘어갑니다.
3. 다른 집합이라면, 각각의 집합이 홍수가 일어날 수 있는지 a, b 배열을 통해서 확인해줍니다.
4. 홍수가 일어날 수 있다면, 미리 구해둔, 홍수가 발생할 수 있는 도시 수에서 집합의 cnt 값을 빼줍니다.
5. 두 집합 중 한쪽에 a, b, cnt를 모두 더해줍니다.
6. 더한 a값이 b값보다 작다면 cnt를 홍수가 발생할 수 있는 도시 수에 더해줍니다.
7. 두 집합을 합칩니다.

2번 쿼리는 전역변수에 저장되어 있는, '홍수가 발생할 수 있는 도시 수'를 출력만 해주면 됩니다.

코드

#include <bits/stdc++.h>

using namespace std;

const int MAX = 100001;
int ans;
int a[MAX], b[MAX], cnt[MAX], parent[MAX];

int Find(int a) {
	if (a == parent[a]) {
		return a;
	}
	else {
		return parent[a] = Find(parent[a]);
	}
}

void Union(int _a, int _b) {
	int pa = Find(_a), pb = Find(_b);

	if (pa != pb) {
		if (pa > pb) {
			swap(pa, pb);
		}

		if (a[pa] < b[pa]) {
			ans -= cnt[pa];
		}

		if (a[pb] < b[pb]) {
			ans -= cnt[pb];
		}

		a[pa] += a[pb];
		b[pa] += b[pb];
		cnt[pa] += cnt[pb];

		if (a[pa] < b[pa]) {
			ans += cnt[pa];
		}

		parent[pb] = pa;

	}
}

int main(void) {
	ios_base::sync_with_stdio(false);
	cin.tie(nullptr);
	cout.tie(nullptr);

	int n, m;
	cin >> n >> m;

	for (int i = 1; i <= n; i++) {
		cin >> a[i];
		cnt[i] = 1;
		parent[i] = i;
	}

	for (int i = 1; i <= n; i++) {
		cin >> b[i];

		if (a[i] < b[i]) {
			ans++;
		}
	}

	while (m--) {
		int q;
		cin >> q;

		if (q == 1) {
			int x, y;
			cin >> x >> y;
			Union(x, y);
		}
		else {
			cout << ans << '\n';
		}
	}

	return 0;
}