P.1647 도시 분할 계획

코드

1. 프림 알고리즘

package Baekjoon.Class;

import java.io.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;

public class P_1647 {

    static int n, m;
    static ArrayList<Node>[] graph;

    static public class Node {
        int end;
        int fee;

        public Node(int end, int fee) {
            this.end = end;
            this.fee = fee;
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        int[] line = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
        n = line[0]; m = line[1];
        graph = new ArrayList[n + 1];
        for (int i = 0; i <= n; i++) graph[i] = new ArrayList<>();
        for (int i = 0; i < m; i++) {
            line = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();

            graph[line[0]].add(new Node(line[1], line[2]));
            graph[line[1]].add(new Node(line[0], line[2]));
        }

        bw.write(Integer.toString(prim(1)));
        bw.flush();

        bw.close();
        br.close();
    }

    private static int prim(int start) {
        PriorityQueue<Node> queue = new PriorityQueue<>(Comparator.comparing(o -> o.fee));
        queue.add(new Node(start, 0));
        boolean[] visited = new boolean[n + 1];
        int total = 0;
        int maxFee = 0;

        while (!queue.isEmpty()) {
            Node poll = queue.poll();

            if (!visited[poll.end]) {
                visited[poll.end] = true;

                total += poll.fee;
                maxFee = Math.max(maxFee, poll.fee);

                for (Node node : graph[poll.end]) {
                    if (!visited[node.end]) {
                        queue.add(node);
                    }
                }
            }
        }

        return total - maxFee;
    }
}

2. 크루스칼 알고리즘

package Baekjoon.Class;

import java.io.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;

public class P_1647_2 {

    static int n, m;
    static ArrayList<Edge> list = new ArrayList<>();
    static int[] parent;

    static public class Edge {
        int start, end;
        int fee;

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

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        int[] line = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
        n = line[0]; m = line[1];
        for (int i = 0; i < m; i++) {
            line = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();

            list.add(new Edge(line[0], line[1], line[2]));
        }
        Collections.sort(list, Comparator.comparing(o -> o.fee));

        parent = new int[n + 1];
        for (int i = 1; i <= n; i++) parent[i] = i;

        bw.write(Integer.toString(kruskal()));
        bw.flush();

        bw.close();
        br.close();
    }

    private static int kruskal() {
        int total = 0;
        int maxFee = 0;

        for (Edge edge : list) {
            if (!isSameParent(edge.start, edge.end)) {
                total += edge.fee;
                maxFee = edge.fee;

                union(edge.start, edge.end);
            }
        }

        return total - maxFee;
    }

    private static void union(int start, int end) {
        start = find(start);
        end = find(end);

        if (start != end) parent[end] = start;
    }

    private static boolean isSameParent(int start, int end) {
        if (find(start) == find(end)) return true;

        return false;
    }

    private static int find(int start) {
        if (parent[start] == start) return start;

        return parent[start] = find(parent[start]);
    }
}

코드 설명

MST와 더불어 전반적인 두 알고리즘에 대한 설명은 다음 글에 있다.

2023.08.17 - [코딩테스트_JAVA/백준] - P.1197 최소 스패닝 트리

P.1197 최소 스패닝 트리

 

프림 알고리즘을 사용할 때 주의해야 할 점은 maxFee를 구할 때 maxFee = poll.fee;로 구하면 안된다는 것이다.

프림 알고리즘은 크루스칼 알고리즘과 다르게 특정 정점에서의 최소 간선을 시작으로 한다.

물론 해당 간선의 가중치가 전체 간선의 가중치 중에서 최솟값일 수 있지만 필연적인 것은 아니므로 Math.max연산을 통해 업데이트 해 주어야 맞는 최댓값을 구할 수 있다.

 

크루스칼 알고리즘은 이전 코드와 유사하게 작성했을 때 시간초과가 발생했다.

이유는 find 함수에 있었다.

private static int find(int edge) {
    if (parent[edge] == edge) return edge;

    return find(parent[edge]);
}

private static int find(int start) {
    if (parent[start] == start) return start;

    return parent[start] = find(parent[start]);
}

첫 번째는 시간 초과가 발생하는 find이고 두 번째는 시간 초과가 발생하지 않은 find다.

두 번째 find는 해당 서브트리에서 루트를 찾아갈 때마다 지속적으로 업데이트를 해 주며 거슬러 올라간다. 트리를 평평하게 만들어주는 셈이다.

이 경우 find 함수의 호출 횟수가 많아졌을 때 굳이 한 칸씩 위로 올라가지 않고 바로 최상단에 접근할 수 있어 더 빠르게 루트를 찾을 수 있다.