[백준][C++] 15809: 전국시대

문제 링크: https://www.acmicpc.net/problem/15809

 

유니온 파인드를 적절히 사용해주면 됩니다.

오랜만에 한번에 맞았네요;; 신기합니다.

 

병력 배열을 만들어 놓고, 두 나라가 동맹을 맺어 합칠 때 병력을 다 root로 합칩니다. 그러면 앞으로 합칠 때 root에 있는 병력만 합치면 됩니다.

전쟁도 동일하게 나라를 합친 뒤 병력의 차를 root로 옮겨놓습니다.

 

'두 나라의 병력이 같을 경우 두 나라 모두 멸망한다'라고 되어 있는데 병력이 0이 되면 멸망하기 때문에 합치든 말든 상관이 없습니다.

 

동맹끼리 다시 동맹을 맺거나 전쟁을 하지 않기 때문에 유니온파인드를 적절히 사용할 수 있습니다.

모든 기록이 끝났을 때는 병력 배열을 체크해서 0이 아닌 개수를 센 다음 오름차순으로 출력하면 됩니다. upper_bound를 적절히 사용하면 되겠죠?

 

#pragma GCC optimize ("Ofast")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2")

#define _CRT_SECURE_NO_WARNINGS
#include <bits/stdc++.h>

using namespace std;


vector<int> parent;
vector<int> soldiers;	//병력

void init(int n) {
	parent.resize(n);
	for (int i = 0; i < n; ++i)
		parent[i] = i;
}

int find(int u) {
	if (u == parent[u]) return u;
	return parent[u] = find(parent[u]);
}

void merge(int u, int v) {
	u = find(u), v = find(v);
	if (u == v) return;

	parent[u] = v;
}


int main()
{
#ifndef ONLINE_JUDGE
	freopen("input.txt", "r", stdin);
	//freopen("output.txt", "w", stdout);
#endif

	ios_base::sync_with_stdio(false);
	cin.tie(nullptr);
	cout.tie(nullptr);

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

	init(n);

	soldiers.resize(n);
	for (auto& c : soldiers)
		cin >> c;

	int o, p, q;
	while (m--) {
		cin >> o >> p >> q;
		--p, --q;
		p = find(p), q = find(q);

		if (o == 1) {
			merge(p, q);

			soldiers[q] += soldiers[p];
			soldiers[p] = 0;
		}
		else {
			if (soldiers[p] > soldiers[q]) swap(p, q);
			
			merge(p, q);

			soldiers[q] -= soldiers[p];
			soldiers[p] = 0;
		}
	}

	sort(soldiers.begin(), soldiers.end());
	auto it = upper_bound(soldiers.begin(), soldiers.end(), 0);
	cout << soldiers.end() - it << '\n';
	
	for (; it != soldiers.end(); ++it)
		cout << *it << ' ';
	cout << '\n';

	return 0;
}

기존 유니온 파인드 코드는 건들지 않았습니다.

전쟁할 때는 병력 많은 쪽으로 합치도록 했고요, 합칠때는 루트쪽으로 병력을 합치도록 했습니다.