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

문제 링크

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

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

문제

상근이는 매일 아침 실험실로 출근해서 샘플의 무게를 재는 일을 하고 있다. 상근이는 두 샘플을 고른 뒤, 저울을 이용해서 무게의 차이를 잰다.

교수님의 마음에 들기 위해서 매일 아침부터 무게를 재고 있지만, 가끔 교수님이 실험실에 들어와서 상근이에게 어떤 두 샘플의 무게의 차이를 물어보기도 한다. 이때, 상근이는 지금까지 잰 결과를 바탕으로 대답을 할 수도 있고, 못 할 수도 있다.

상근이는 결과를 출근한 첫 날부터 공책에 적어 두었다. 하지만, 엄청난 양의 무게가 적혀있기 때문에, 교수님의 질문에 재빨리 대답할 수가 없었다. 이런 상근이를 위해서 프로그램을 만들려고 한다.

상근이가 실험실에서 한 일이 순서대로 주어진다. 어떤 두 샘플의 무게의 차이를 구할 수 있는지 없는지를 알아내는 프로그램을 작성하시오.

입력

입력은 여러 개의 테스트 케이스로 이루어져 있다.

각 테스트 케이스의 첫째 줄에는 샘플의 종류의 개수 N (2 ≤ N ≤ 100,000)과 상근이가 실험실에서 한 일의 수 M (1 ≤ M ≤ 100,000)이 주어진다. 샘플은 1번부터 N번까지 번호가 매겨져 있다. 다음 M개 줄에는 상근이가 실험실에서 한 일이 주어진다.

상근이가 무게를 쟀다면, ! a b w와 같은 형식으로 주어진다. 이 뜻은 b가 a보다 w그램 무겁다는 뜻이다. (a ≠ b) w는 1,000,000을 넘지 않는 음이 아닌 정수이다. 모든 측정은 정확하고, 일관성을 유지한다.

교수님의 질문은 ? a b와 같은 형식으로 주어진다. 이 뜻은 b가 a보다 얼마나 무거운지를 출력하라는 뜻이다.

마지막 줄에는 0이 두 개 주어진다.

출력

교수님의 질문 (? a b)이 입력으로 들어올 때 마다, 지금까지 측정한 결과를 바탕으로 a와 b의 무게 차이를 계산할 수 있다면, b가 a보다 얼마나 무거운지를 출력한다. 무게의 차이의 절댓값이 1,000,000을 넘지 않는다. 만약, 측정한 결과를 바탕으로 무게의 차이를 계산할 수 없다면, "UNKNOWN"을 출력한다.

예제 입력 1

2 2
! 1 2 1
? 1 2
2 2
! 1 2 1
? 2 1
4 7
! 1 2 100
? 2 3
! 2 3 100
? 2 3
? 1 3
! 4 3 150
? 4 1
0 0

예제 출력 1

1
-1
UNKNOWN
100
200
-50

알고리즘 분류

• 유니온 파인드

풀이

이 형님이 쓰신 글을 참고하였다.

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

명령어가 !이면 A, B를 유니온 파인드로 묶어준다. 이 때 B가 A보다 W만큼 무거우므로, A를 B의 Parent로 정해주고 A에 대한 B의 상대적인 무게를 Dist[B]에 기록해준다. 따라서 Dist[B]의 의미는 B의 무게는 B의 Parent인 A보다 Dist[B]만큼 무겁다라는 의미이다.

코드

#include <iostream>
#include <string>
#include <cmath>
#include <set>
#include <map>
#include <stack>
#include <queue>
#include <vector>
#include <algorithm>
#include <climits>
#define FASTIO cin.tie(NULL); cout.tie(NULL); ios::sync_with_stdio(false);
#define MAX 100001
#define LL long long
#define INF 1e9

using namespace std;
int N, M;
int Parent[MAX];
int Dist[MAX];

void Init() {
	for (int i = 1; i <= N; i++) {
		Parent[i] = i;
		Dist[i] = 0;
	}
}

int Find(int X) {
	if (Parent[X] == X) {
		return X;
	}
	int PA = Find(Parent[X]);
	Dist[X] += Dist[Parent[X]]; // Dist[Node] = (Node부터 Node의 Parent까지의 거리) + (Node의 Parent부터 Root까지의 거리)
	// Dist[Node]를 갱신하려면, Node의 Parent부터 Root까지의 거리를 더해준다.
	return Parent[X] = PA;
}

void Union(int X, int Y, int W) {
	int PX = Find(X);
	int PY = Find(Y);
	if (PX != PY) { // Y가 X보다 W만큼 더 무거우므로, Y의 Parent를 X로 기록해둔다.
		Dist[PY] = (-Dist[Y]) + Dist[X] + W; // 그리고 Parent와 Node의 상대적 무게를 기록한다.
		// Y의 무게는 Y의 Parent와 -Dist[Y]만큼의 무게 차이가 난다.
		Parent[PY] = PX;
	}
}

void Query() {
	while (M--) {
		char CMD;
		cin >> CMD;
		if (CMD == '!') {
			int A, B, W;
			cin >> A >> B >> W;
			Union(A, B, W);
		}
		else if (CMD == '?') {
			int A, B;
			cin >> A >> B;
			int PA = Find(A);
			int PB = Find(B);
			if (PA == PB) {
				cout << Dist[B] - Dist[A] << "\n";
			}
			else {
				cout << "UNKNOWN\n";
			}
		}
	};
}

void Input() {
	while (1) {
		cin >> N >> M;
		if ((N == 0) && (M == 0)) {
			break;
		}
		Init();
		Query();
	};
}

int main() {
	FASTIO
	Input();

	return 0;
}