[BOJ 3830] 교수님은 기다리지 않는다

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

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

 

- 문제 요약

 

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

 

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

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

 

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

 

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

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

 

 

- 알고리즘 정리

 

Union-Find 알고리즘을 사용해서 풀었습니다.

기본적으로 Union-Find에서 사용하는 Union 함수, Find 함수, parent 배열 등을 변형해봤습니다.

arr 배열에 루트 노드까지의 거리를 저장해주고, Find 함수에서 루트 노드를 갱신시켜줍니다. Union 함수에는 만들어놨던 arr 배열로 집합을 만들어줍니다.

 

 

- 코드 작성

 

#include<bits/stdc++.h>
using namespace std;

#define MAX 100005
typedef long long ll;
int n,m,p[MAX];
ll arr[MAX];

int find(int node){
	if(p[node]==-1){
		return node;
	}
	int parent=find(p[node]);
	arr[node]+=arr[p[node]];
	return p[node]=parent;
}

void Union(int a,int b,int w){
	int fa=find(a);
	int fb=find(b);
	if(fa==fb){
		return;
	}
	arr[fb]=arr[a]-arr[b]+w;
	p[fb]=fa;
	return;
}

int main(){
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);
	while(true){
		memset(p,-1,sizeof(p));
		memset(arr,0,sizeof(arr));
		cin>>n>>m;
		if(n==0&&m==0){
			break;
		}
		char c;
		int a,b,w;
		for(int i=0;i<m;i++){
			cin>>c;
			if(c=='!'){
				cin>>a>>b>>w;
				Union(a,b,w);
			}
			else{
				cin>>a>>b;
				if(find(a)!=find(b)){
					cout<<"UNKNOWN\n";
				}
				else{
					cout<<arr[b]-arr[a]<<'\n';
				}
			}
		}
	}
}

코드 제출 결과