[PS] 교수님은 기다리지 않는다 / C++ (백준 3830번)

교수님은 기다리지 않는다 (백준 3830번)

 

분리 집합을 떠올리는데는 큰 어려움은 없었지만 가중치 관리하는 수식을 세우기가 조금 힘들었던 문제

 

풀이 과정

 분리 집합 알고리즘을 구현한 뒤 경로 압축을 해준다(find_root() 함수). _union() 함수에 알맞은 수식을 넣어 가중치를 갱신하게끔 하면 된다.

 

find_root(node): node의 최종 root를 찾은 뒤, 경로 압축으로 최적화 실행

_union(a, b, w): a에서 b로 가는 가중치가 w라고 했을 때, a의 최종 root(pa_a)와 b의 최종 root(pa_b)를 각각 구하여 가중치를 갱신

parent[pa_a.first] = pa_b.first; // a의 최종 root(pa_a.first)의 최종 root(parent[pa_a.first])를 b의 최종 root(pa_b.first)로 설정

val[pa_a.first] = pa_b.second + w - pa_a.second; // (b의 최종 root가 b보다 얼마나 무거운지) + (b가 a보다 얼마나 무거운지) - (a의 최종 root가 a보다 얼마나 무거운지)

 

#include<iostream>
#include<algorithm>
#include<cstring>
using namespace std;
typedef long long ll;

int parent[100001];
ll val[100001];

pair<int, ll> find_root(int node) {
	if (parent[node] == 0)return make_pair(node, 0);
	auto pa = find_root(parent[node]);
	val[node] += pa.second;
	parent[node] = pa.first;
	return make_pair(pa.first, val[node]);
}

void _union(int a, int b, ll w) {
	auto pa_a = find_root(a);
	auto pa_b = find_root(b);
	if (pa_a.first == pa_b.first) return;
	parent[pa_a.first] = pa_b.first;
	val[pa_a.first] = pa_b.second + w - pa_a.second;
	return;
}

bool logic()
{
	int n, m;
	scanf("%d %d", &n, &m);
	if (n == 0 && m == 0)return false;
	for (int i = 0; i < m; i++) {
		char ch;
		int a, b;
		scanf(" %c %d %d", &ch, &a, &b);
		if (ch == '?') {
			auto pa_a = find_root(a);
			auto pa_b = find_root(b);
			if (pa_a.first != pa_b.first)printf("UNKNOWN\n");
			else printf("%d\n", pa_a.second - pa_b.second);

		}
		else {
			ll w;
			scanf("%lld", &w);
			_union(a, b, w);
		}
	}
	return true;
}

int main()
{
	while (logic()) {
		memset(parent, 0, sizeof(parent));
		memset(val, 0, sizeof(val));
	}
	return 0;
}