[알고리즘] Java / 백준 / 전기가 부족해 / 10423

[알고리즘] Java / 백준 / 전기가 부족해 / 10423

문제

문제 링크

접근 방식

크루스칼 알고리즘으로 사이클이 생기지 않으면서 최소 가중치의 간선을 선택하되 유니온 파인드의 루트를 발전소의 위치로 하여 한 정점이 2개 이상의 발전소와 이어지지 않도록 한다.

간선을 연결할 때마다 모든 정점이 발전소와 이어져 있는지 확인하고, 모두 발전소와 이어져 있다면 간선 잇는 것을 멈춘다.

코드

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

public class Main_10423 {
	static int[][] parent;
	static int N,M,K;
	static List<Integer> powers;
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		N = Integer.parseInt(st.nextToken());
		M = Integer.parseInt(st.nextToken());
		K = Integer.parseInt(st.nextToken());
		
		powers = new ArrayList<>();
		st = new StringTokenizer(br.readLine());
		for(int i=1;i<=K;i++) {
			powers.add(Integer.parseInt(st.nextToken()));
		}
		
		int[][] graph = new int[M+1][3];
		
		for(int i=1;i<=M;i++) {
			st = new StringTokenizer(br.readLine());
			int u = Integer.parseInt(st.nextToken());
			int v = Integer.parseInt(st.nextToken());
			int w = Integer.parseInt(st.nextToken());
			
			graph[i] = new int[] {u,v,w};
		}
		
		Arrays.sort(graph,(o1,o2)->Integer.compare(o1[2], o2[2]));
		
		parent = new int[N+1][2];
		makeSet();
		
		int cost = 0;
		for(int i=1;i<=M;i++) {
			if(union(graph[i][0],graph[i][1])) {
				cost += graph[i][2];
			}
			
			if(connect()) break;
		}
		
		System.out.println(cost);
		
	}
	
	public static boolean connect() {
		boolean isConnect = true;
		
		List<Integer> powerRoot = new ArrayList<>();
		
		for(int i=0;i<K;i++) {
			powerRoot.add(find(powers.get(i)));
		}
		
		
		for(int i=1;i<=N;i++) {
			if(powerRoot.contains(find(i))) continue;
			isConnect = false;
			break;
		}
		
		return isConnect;
	}
	
	public static void makeSet() {
		for(int i=1;i<=N;i++) {
			if(powers.contains(i))parent[i][1] = 1;
			parent[i][0] = i;
		}
	}
	
	public static boolean union(int a, int b) {
		
		
		a = find(a);
		b = find(b);
		
		if(a!=b && parent[a][1] + parent[b][1] != 2) {
			if(parent[a][1] == 1) {
				parent[b][0] = a;
			}else if(parent[b][1] == 1){
				parent[a][0] = b;
			}else {
				parent[b][0] = a;				
			}
			
			return true;
		}
		return false;
	}
	
	public static int find(int a) {
		if(parent[a][0] == a) return a;
		
		return parent[a][0] = find(parent[a][0]);
	}

}