[BOJ 16562] 친구비 (Java)

BOJ 16562 친구비

문제풀이

Disjoint Set을 이용하는 문제다. 처음에는 Union-find 함수를 통해서 모두 그룹지어놓고 백트래킹을 통해 완전탐색을 생각했다. 하지만 시간초과하였고 다시 생각해보니 전혀 완전탐색을 할 필요가 없었다.

1. 주어진 친구 관계를 바탕으로 Disjoint Set으로 만든다.
2. 각 Set별로 친구비가 가장 적은 친구를 찾는다.
3. 해당 친구들과 모두 친구를 맺고 전체 친구비가 예산안에 들어오는지 확인한다.

• Union-find를 사용할 때 주의할 점은 parent 배열을 직접 접근하면 안되고 항상 find(하지만 ) 함수를 통해서 접근해야하는 것이다.
• 트리의 높이를 저장하고 있으면 항상 높이가 낮은 트리를 높은 트리에 붙여서 최적화 할 수 있다.

구현코드

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

public class Main {
	
	static int[] rank;
	static int[] parent;
	static int[] cost;
	static int N, M, K, ans;
	
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		N = stoi(st.nextToken());
		M = stoi(st.nextToken());
		K = stoi(st.nextToken());
	
		ans = 0;
		rank = new int[N + 1];
		parent = new int[N + 1];
		cost = new int[N + 1];
		
		for(int i = 0 ; i <= N ; ++i) {
			parent[i] = i;
		}
		
		st = new StringTokenizer(br.readLine());
		for(int i = 1 ; i <= N ; ++i) {
			cost[i] = stoi(st.nextToken());
		}
		
		for(int i = 0 ; i < M ; ++i) {
			st = new StringTokenizer(br.readLine());
			
			int p1 = stoi(st.nextToken());
			int p2 = stoi(st.nextToken());
			if(find(p1) != find(p2)) {
				union(p1, p2);
			}
		}
		
		int[] min = new int[N + 1];
		Arrays.fill(min, Integer.MAX_VALUE);
		
		for(int i = 1 ; i <= N ; ++i) {
			int group = find(i);
			min[group] = min[group] > cost[i] ? cost[i] : min[group];
		}
		
		for(int i = 1 ; i <= N ; ++i) {
			if(min[i] != Integer.MAX_VALUE) {
				ans += min[i];
			}
		}
		
		if(ans > K) System.out.println("Oh no");
		else System.out.println(ans);
	}

	private static void union(int a, int b) {
		int rootA = find(a);
		int rootB = find(b);
		
		if(rank[rootA] < rank[rootB]) {
			parent[rootA] = rootB;
		} else {
			parent[rootB] = rootA;
			
			if(rank[rootA] == rank[rootB]) {
				rank[rootA]++;
			}
		}
	}
	
	private static int find(int a) {
		if(parent[a] == a) return a;
		
		return parent[a] = find(parent[a]);
	}
	
	private static int stoi(String s) {
		return Integer.parseInt(s);
	}
}