[알고리즘] 백준 10423 전기가 부족해 -최소신장트리, 크루스칼- 자바 코틀린

 

 

www.acmicpc.net/problem/10423

10423번: 전기가 부족해

 

 

 

최소신장트리 유형 문제를 풀어봤습니다. ㅎㅎ 오타하나떄문에 시간을 많이 날렸습니다. ㅠㅠ ;;

 

문제를 간단히 설명하면, 발전소에서 전기가 나오고 도시들이 모두 전기를 저렴하게 연결하는 비용을 구하면 됩니다.

크루스칼 알고리즘을 사용하여 풀면됩니다. 

 

기본적인 최소신장트리 문제와 다른게 있다면, 모든 도시들을 그냥 싼 비용으로 연결하는게 아니라 여러개의 발전소가 있고 발전소로부터 전기를 공급만 다 받으면 되어서 하나의 연결집합이 아니라 여러개의 연결집합이 나올 수 있습니다. 

 

 

그래서 발전소를 powerPantList 에 담아 Union-find 함수들을 알맞게 수정했습니다.

private static boolean isSameParent(int a, int b) {
        if (powerPantList.contains(find(a)) && powerPantList.contains(find(b))) return true; // 둘다 발전소에 연결된 경우 연결된걸로 취급
        return find(a) == find(b);
    }

연결되어 있는건지 판단하는 isSameParent() 에서 두 집합이 서로 연결이 안되어있어도 각각 발전소에 연결된거면 서로 연결할 필요가 없으니 연결된걸로 판단한다.

 

private static void union(int a, int b) {
        a = find(a);
        b = find(b);
        if (!powerPantList.contains(a) && powerPantList.contains(b)) { // b만 발전소에 연결되어 있는 경우
            parents[a] = b;
        } else if (powerPantList.contains(a) && !powerPantList.contains(b)) { // a만 발전소에 연결되어 있는 경우
            parents[b] = a;
        } else { // 둘 다 아직 발전소에 연결 안된 경우
            if (a < b) {
                parents[b] = a;
            } else {
                parents[a] = b;
            }
        }
    }

union() 할 시 둘 다 발전소에 연결 안된 경우는 두 집합을 연결해주고 하나의 집합만 발전소에 연결되어 있는 경우 전기를 공급받기 위해 연결된 집합으로 Union 해준다. 

 

-

풀면서 그냥 쓴 메모..

 

풀이는 주석으로 자세히 적어놔서 생략하겠습니다.

 

[Java]

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Scanner;

public class Main {
    private static int[] parents;
    private static List<Edge> edgeList = new ArrayList<>();
    private static int N; // 도시의 개수
    private static int M; // 케이블의 개수
    private static int K; // 발전소의 개수
    private static List<Integer> powerPantList = new ArrayList<>(); // 발전소 번호 리스트
    private static int answer = 0;


    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        N = sc.nextInt();
        M = sc.nextInt();
        K = sc.nextInt();
        parents = new int[N + 1];
        for (int i = 1; i <= N; i++) {
            parents[i] = i;
        }
        //발전소가 설치될 도시의 번호
        for (int i = 0; i < K; i++) {
            powerPantList.add(sc.nextInt());
        }
        // 엣지리스트 초기화
        for (int i = 0; i < M; i++) {
            int start = sc.nextInt();
            int end = sc.nextInt();
            int weight = sc.nextInt();
            edgeList.add(new Edge(start, end, weight));
        }
        // 크루스칼 알고리즘
        Collections.sort(edgeList);
        for (Edge edge : edgeList) {
            if (!isSameParent(edge.start, edge.end)) { //서로 연결 안된경우 union
                answer += edge.weight;
                union(edge.start, edge.end);
            }
        }

        System.out.println(answer);
    }

    private static void union(int a, int b) {
        a = find(a);
        b = find(b);
        if (!powerPantList.contains(a) && powerPantList.contains(b)) { // b만 발전소에 연결되어 있는 경우
            parents[a] = b;
        } else if (powerPantList.contains(a) && !powerPantList.contains(b)) { // a만 발전소에 연결되어 있는 경우
            parents[b] = a;
        } else { // 둘 다 아직 발전소에 연결 안된 경우
            if (a < b) {
                parents[b] = a;
            } else {
                parents[a] = b;
            }
        }
    }

    private static int find(int a) {
        if (parents[a] == a) {
            return a;
        } else {
            return parents[a] = find(parents[a]);
        }
    }

    private static boolean isSameParent(int a, int b) {
        if (powerPantList.contains(find(a)) && powerPantList.contains(find(b))) return true; // 둘다 발전소에 연결된 경우 연결된걸로 취급
        return find(a) == find(b);
    }


    static class Edge implements Comparable<Edge> {
        int start;
        int end;
        int weight;

        Edge(int start, int end, int weight) {
            this.start = start;
            this.end = end;
            this.weight = weight;
        }

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

 

[Kotlin]

import java.util.*

private lateinit var parents: IntArray
private val edgeList: MutableList<Edge> = ArrayList()
private var N // 도시의 개수
        = 0
private var M // 케이블의 개수
        = 0
private var K // 발전소의 개수
        = 0
private val powerPantList: MutableList<Int> = ArrayList() // 발전소 번호 리스트

private var answer = 0


fun main(args: Array<String>) {
    val sc = Scanner(System.`in`)
    N = sc.nextInt()
    M = sc.nextInt()
    K = sc.nextInt()
    parents = IntArray(N + 1)
    for (i in 1..N) {
        parents[i] = i
    }
    //발전소가 설치될 도시의 번호
    for (i in 0 until K) {
        powerPantList.add(sc.nextInt())
    }
    // 엣지리스트 초기화
    for (i in 0 until M) {
        val start = sc.nextInt()
        val end = sc.nextInt()
        val weight = sc.nextInt()
        edgeList.add(Edge(start, end, weight))
    }
    // 크루스칼 알고리즘
    edgeList.sort()
    for (edge in edgeList) {
        if (!isSameParent(edge.start, edge.end)) { //서로 연결 안된경우 union
            answer += edge.weight
            union(edge.start, edge.end)
        }
    }
    println(answer)
}

fun union(a: Int, b: Int) {
    var a = a
    var b = b
    a = find(a)
    b = find(b)
    if (!powerPantList.contains(a) && powerPantList.contains(b)) { // b만 발전소에 연결되어 있는 경우
        parents[a] = b
    } else if (powerPantList.contains(a) && !powerPantList.contains(b)) { // a만 발전소에 연결되어 있는 경우
        parents[b] = a
    } else { // 둘 다 아직 발전소에 연결 안된 경우
        if (a < b) {
            parents[b] = a
        } else {
            parents[a] = b
        }
    }
}

fun find(a: Int): Int {
    return if (parents[a] == a) {
        a
    } else {
        find(parents[a]).also { parents[a] = it }
    }
}

fun isSameParent(a: Int, b: Int): Boolean {
    return if (powerPantList.contains(find(a)) && powerPantList.contains(find(b))) true else find(a) == find(b) // 둘다 발전소에 연결된 경우 연결된걸로 취급
}


class Edge(var start: Int, var end: Int, var weight: Int) : Comparable<Edge?> {
    override operator fun compareTo(other: Edge?): Int {
        return weight - other!!.weight
    }

}

 

댓글과 공감은 큰 힘이 됩니다. 감사합니다. !!!