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

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

문제 요약

입력은 여러 개의 테스트 케이스로 주어진다. 끝날 때는 두 개의 0이 주어진다.
N개의 샘플 종류와 M번의 작업 개수가 주어진다.
"! a b w"는 b가 a보다 w그램 무겁다는 뜻이다.
"? a b"는 b가 a보다 얼마나 무거운지 출력하는 것이다. 얻을 수 없다면 "UNKNOWN"을 출력한다.

N은 최대 10만, M은 최대 10만, w는 100만 이하의 음이 아닌 정수이다.
무게 차이의 절대값은 100만을 넘지 않는다 -> int 사용

---

문제 접근

유니온 파인드(Union-Find)를 기반으로 문제를 해결할 수 있다.
다만 집합을 나타낼 때, 현재 노드와 루트와의 무게 차이도 기록해야 한다.
무거운 샘플이 루트 쪽으로 가도록 한다. (사실 갱신하다보면 경중이 뒤집힐 수는 있다)
무게 차이는 "현재 노드 무게 + 무게 차이 = 루트 무게" 가 되도록 기록하자.
샘플의 번호가 1부터 N까지로 사용하니까 크기를 N+1로 선언하자.

두 샘플의 차이를 비교할 때, 둘의 부모를 따라가면서 그 차이를 더하고 빼는 방식으로 계산하면 시간초과가 난다.
find를 할 때 지나는 노드가 모두 루트를 가리키게 하면서 무게도 갱신시키면 해결할 수 있다.

---

풀이

기본적인 아이디어는 다음과 같다.

1. (정리 예정)

이를 구현한 코드는 다음과 같다.

import java.util.*;
import java.io.*;

public class Main {
	static int[] parent;
	static int[] weightDiff;
	
	// 해당 노드의 루트 얻기
	// a + weightDiff[a] = parent[a]
	// parent[a] + weightDiff[parent[a]] = parent[parent[a]]
	// a + weightDiff[a] + weightDiff[parent[a]] + ... = aRoot;
	public static int[] find(int a) {
		if (parent[a] == a) return new int[] {a, weightDiff[a]};
		else {
			int[] now = find(parent[a]); 
			parent[a] = now[0];
			weightDiff[a] += now[1];
			now[1] = weightDiff[a];
			
			return now;
		}
	}
	
	// a의 루트가 b의 루트를 가리키도록 만든다.
	// a + w = b
	// a + weightDiff[a] = aRoot
	// b + weightDiff[b] = bRoot 이므로,
	
	// aRoot - weightDiff[a] + w = bRoot - weightDiff[b]
	// aRoot - weightDiff[a] + weightDiff[b] + w = bRoot
	public static void union(int a, int b, int w) {
		int aRoot = find(a)[0];
		int bRoot = find(b)[0];
		
		parent[aRoot] = bRoot;
		weightDiff[aRoot] = w - weightDiff[a] + weightDiff[b];
	}
	
    public static void main(String[] args) throws IOException {
    	BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    	StringBuilder sb = new StringBuilder();
    	
    	StringTokenizer st1 = new StringTokenizer(br.readLine());
    	
    	int N = Integer.parseInt(st1.nextToken());
    	int M = Integer.parseInt(st1.nextToken());
    	
    	// 테스트 케이스
    	while (N > 0 && M > 0) {
    		
    		// 필요한 배열 선언 및 초기화
    		// parent는 해당 노드의 부모(또는 루트)를 가리키며, 그 무게 차이를 weightDiff에 기록한다.
    		parent = new int[N+1];
    		weightDiff = new int[N+1];
    		
    		for (int i = 0; i < N+1; i++) {
    			parent[i] = i;
    		}
    		for (int i = 0; i < N+1; i++) {
    			weightDiff[i] = 0;
    		}
    		
    		// M번의 명령을 받으며 결과를 기록한다.
    		for (int i = 0; i < M; i++) {
    			
    			StringTokenizer stm = new StringTokenizer(br.readLine());
    			
    			String order = stm.nextToken();
    			// 무게를 측정하는 경우
    			if (order.equals("!")) {
    				int a = Integer.parseInt(stm.nextToken());
    				int b = Integer.parseInt(stm.nextToken());
    				int w = Integer.parseInt(stm.nextToken());
    				
    				union(a, b, w);
    			}
    			// 무게를 비교하는 경우
    			else if (order.equals("?")) {
    				int a = Integer.parseInt(stm.nextToken());
    				int b = Integer.parseInt(stm.nextToken());
    				
    				// 루트가 같은지 확인해서 만약 둘이 같은 집합에 있다면 무게 차이 계산해서 출력에 기록
    				// b가 a보다 얼마나 무거운지?
    				
    				// a + weightDiff[a] = aRoot;
    				// b + weightDiff[b] = bRoot;
    				// aRoot = bRoot
    				
    				// a + weightDiff[a] - weightDiff[b] = b
    				
    				if (find(a)[0] == find(b)[0]) {
    					int diff = weightDiff[a] - weightDiff[b];
    					sb.append(diff+"\n");
    				}
    				// 같은 집합에 있지 않다면 무게 차이 계산 불가능
    				else sb.append("UNKNOWN\n");
    			}
    		}
    		
    		// 다음 테스트 케이스의 N과 M을 입력받음
    		// 둘 다 0이라면 테스트 케이스 끝
    		StringTokenizer stNext = new StringTokenizer(br.readLine());
    		N = Integer.parseInt(stNext.nextToken());
    		M = Integer.parseInt(stNext.nextToken());
    	}

    	// 결과 출력
    	System.out.println(sb);
    }
}

---

회고

틀렸던 이유 1
두 샘플의 차이를 비교할 때, 그때 부모를 따라가면서 차이를 더하고 빼는 방식으로 수행해서 시간초과가 발생했다.