[백준] 10423: 전기가 부족해 JAVA

접근

기본적으로는 크루스칼 알고리즘으로 푸는데, 이제 같은 트리라는 것만 확안하는 것이 아니라 한 트리 안에 발전소가 1개만 있어야 한다는 조건이 추가된다. 그래서 union을 할 때에는 루트값을 무조건 발전소가 있는 번호가 저장되도록 설정했고, find로 찾은 루트가 모두 발전소라면 넘어가도록 처리해주었다. 

 

그리고 모든 도시가 하나의 발전소에 연결되기 위해서 필요한 케이블의 개수는 n-k개이다. 설치된 케이블의 개수를 확인하여 n-k가 되면 while문을 종료하도록 했다. 

 

코드

package week7;

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

public class BJ10423 {
    static class Edge implements Comparable<Edge> {
        int from, to, cost;

        public Edge(int from, int to, int cost) {
            this.from = from;
            this.to = to;
            this.cost = cost;
        }

        @Override
        public int compareTo(Edge o) {
            return this.cost - o.cost;
        }
    }

    static int[] root;
    static ArrayList<Integer> powerPlant;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = null;

        st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());

        // 발전소가 설치된 도시의 번호
        powerPlant = new ArrayList<>();
        PriorityQueue<Edge> pq = new PriorityQueue<>();

        st = new StringTokenizer(br.readLine());
        while(st.hasMoreTokens()){
            int val = Integer.parseInt(st.nextToken());
            powerPlant.add(val);
        }

        for (int i = 0; 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());
            pq.add(new Edge(u, v, w));
        }

        // 크루스칼 알고리즘
        root = new int[n+1];
        for (int i = 0; i <= n; i++)
            root[i] = i;

        int answer = 0;
        int cnt = 0;

        while(!pq.isEmpty()){
            if(cnt == n-k) break;

            Edge e = pq.poll();

            int r1 = find(e.from);
            int r2 = find(e.to);

            // 같은 트리 안에 있거나, 둘 다 발전소에 연결되어 있을 경우 넘어감
            if(r1 == r2 || (powerPlant.contains(r1) && powerPlant.contains(r2))) continue;

            union(r1, r2);
            answer += e.cost;
            cnt++;
        }

        System.out.println(answer);
    }

    static int find(int n){
        if(root[n] == n) return n;
        return root[n] = find(root[n]);
    }

    static void union(int r1, int r2){
        if(powerPlant.contains(r1)){
            int tmp = r1;
            r1 = r2;
            r2 = tmp;
        }

        root[r1] = r2;
    }
}