[백준/BOJ] 백준 3830번 : 교수님은 기다리지 않는다

www.acmicpc.net/problem/3830

3830번: 교수님은 기다리지 않는다

유니온 파인드를 활용하여 문제를 해결했다. vector<int> root_gap(100001); 에 자신이 속한 루트와의 무게 차이를 저장했다. (예: root_gap[a] = 100 이면 루트가 a보다 100만큼 더 무거운 것) b가 a보다 w그램 무겁다고 할 때, b 쪽으로 merge 하며 b의 무게 = a무게 + w 라는것을 고려해서 root_gap[a_root] = root_gap[b] + w - root_gap[a] 이렇게 갱신했다. 그리고 Find하여 루트를 찾을 때 루트와의 차이(root_gap)를 업데이트 한다 (루트가 바뀌면 업데이트된다). 그리고 b가 a보다 얼마나 무거운지 확인할 때 b무게 + root_gap[b] = root의 무게, a무게 + root_gap[a] = root의 무게, b무게 - a무게 + root_gap[b] - root_gap[a] = 0, b무게 - a무게 = - root_gap[b] + root_gap[a] 를 통해 root_gap[a] - root_gap[b] 값을 구한다.

 

코드

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

int n, m;
vector<int> parent(100001);
vector<int> root_gap(100001); //예) root_gap[a] = 100 이면 루트 노드가 a보다 100만큼 더 무거운것

void Pre()
{
	for (int i = 0; i < 100001; i++)
	{
		parent[i] = i;
		root_gap[i] = 0;
	}
}

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

	int root_u = Find(parent[u]);
	root_gap[u] += root_gap[parent[u]]; //루트와의 차이 업데이트 (루트가 바뀌면 업데이트 된다)
	return parent[u] = root_u;
}

//b가 a보다 w그램 무거울 때
void Merge(int a, int b, int w)
{
	int a_root = Find(a);
	int b_root = Find(b);

	parent[a_root] = b_root;
	root_gap[a_root] = root_gap[b] + w - root_gap[a]; //b의 무게 = a무게 + w 라는것을 고려해서 계산 
}

int main()
{
	cin.tie(NULL);
	ios_base::sync_with_stdio(false);

	while (1)
	{
		Pre();
		cin >> n >> m;

		if (n == 0 && m == 0)
			break;

		for (int i = 0; i < m; i++)
		{
			char oper;
			int a, b, w;

			cin >> oper;

			if (oper == '!')
			{
				cin >> a >> b >> w; //b가 a보다 w그램 무겁다

				if (Find(a) != Find(b))
				{
					Merge(a, b, w);
				}
			}

			else if (oper == '?')
			{
				cin >> a >> b; //b가 a보다 얼마나 무거운가

				//루트가 다를때
				if (Find(a) != Find(b))
					cout << "UNKNOWN" << "\n";

				else
				{
					//b무게 + root_gap[b] = root의 무게
					//a무게 + root_gap[a] = root의 무게
					//b무게 - a무게 + root_gap[b] - root_gap[a] = 0
					//b무게 - a무게 = - root_gap[b] + root_gap[a]
					cout << root_gap[a] - root_gap[b] << "\n";
				}
			}
		}
	}


	return 0;
}