[코테] 백준 3830번 : 교수님은 기다리지 않는다 (java)

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

알고리즘 분류 : 자료 구조, 분리 집합

❓문제

---

🔅해석

Union-Find 알고리즘으로 해결한다.

 

부모와 가중치를 함께 저장하기 위해 Node class를 정의한다.

 

// TestCase3
4 7
! 1 2 100
? 2 3
! 2 3 100
? 2 3
? 1 3
! 4 3 150
? 4 1

 

1. Initialize

 

2. "! 1 2 100"

 

find를 통해 서로의 부모가 같은지 확인하고 같지 않다면, a의 부모를 b의 부모로 변경하여 하나로 합친다.

 

3. "? 2 3"

2과 3의 부모를 확인한다. 2의 parent(=2) != 3의 parent(=3) 이므로 "UNKNOWN"를 반환한다.

 

4. " ! 2 3 100"

 

 

5. "? 2 3"

2과 3의 부모를 확인한다. 2의 parent(=3) == 3의 parent(=3) 이므로 둘의 가중치 차이를 반환한다.

 

6. "? 1 3"

1과 3의 부모를 확인한다. 1의 parent(=2) != 3의 parent(=3) 이다. 하지만 2의 parent가 3이므로 1과 3은 같은 집합에 있다. 

1에서 2, 2에서 3으로 가는 가중치의 합을 더한다.

 

 

7. "! 4 3 150"

 

8. "? 4 1"

4와 1의 부모를 확인한다. 4의 parent(=3) == 1의 parent(=3) 이므로 둘의 가중치 차이를 반환한다.

---

⭕정답 코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {
	
	static int N, M;
	static Node[] parent;
	static StringBuilder sb = new StringBuilder();

	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st;
		
		String oper = "";
		int a, b, w;
		while (true) {
			st = new StringTokenizer(br.readLine());
			N = Integer.parseInt(st.nextToken());
			M = Integer.parseInt(st.nextToken());
			
			parent = new Node[N + 1];
			for (int n = 1; n <= N; n++) {
				parent[n] = new Node(n, 0);
			}
			
			if (N == 0 && M == 0) break;
			
			for (int m = 0; m < M; m++) {
				
				st = new StringTokenizer(br.readLine());
				oper = st.nextToken();
				a = Integer.parseInt(st.nextToken());
				b = Integer.parseInt(st.nextToken());
				
				if ("!".equals(oper)) {
					w = Integer.parseInt(st.nextToken());
					
					union(a, b, w);
				} else {
					int rootA = find(a);
					int rootB = find(b);
					
					if (rootA != rootB) {
						sb.append("UNKNOWN").append("\n");
					} else {
						sb.append(parent[a].w - parent[b].w).append("\n");
					}
				}
			}
		}
		
		System.out.println(sb.toString());;
		br.close();
	}


	private static int find(int a) {
		if (parent[a].p == a) return a;
		int root = find(parent[a].p);
		parent[a].w += parent[parent[a].p].w;
		return parent[a].p = root;
	}


	private static void union(int a, int b, int w) {
		int rootA = find(a);
		int rootB = find(b);
		
		if (rootA != rootB) {
			parent[rootA].p = rootB;
			parent[rootA].w = parent[b].w + w - parent[a].w;
		}
	}

}

class Node {
	int p;
	int w;
	
	public Node(int p, int w) {
		this.p = p;
		this.w = w;
	}
}

 

---

❗결과