자바 - 백준 1647 / 도시 분할 계획

https://www.acmicpc.net/problem/1647

골드4

구현 방법

최적화 전

"임의의 두 집 사이에 경로가 항상 존재", "비용 양수" 두가지를 통해 최소 스패닝 트리 문제임을 눈치 챘습니다.

여기서 추가 조건으로 두개의 그룹을 만들라고 했습니다. 이는 하나의 최소 스패닝 트리를 만들고,

그 트리의 비용들 중 가장 큰 비용를 끊어주면 두 그룹이 나오고, 각각 최소 비용이 된다고 생각했습니다.

최소 스패닝 트리는 크루스칼 알고리즘을 통해 구현했습니다.

 

하지만 꽤 느렸습니다. 따라서 두가지 방법을 통해 최적화 진행 했습니다.

 

방법1. 트리의 특징을 통한 최적화

트리는 (노드 수 - 1)개의 간선을 가지고 있습니다. 현재 문제에서는 하나의 트리에서 간선 하나를 더 끊어서 두 그룹으로 만들면 됩니다.

따라서 (노드 수 - 2)개의 간선만 선택하도록 구현하면 됩니다.

이렇게 하면 기존 방식은 간선 선택을 할 때마다 뽑힌 간선들 중 가장 비용이 높은 값을 비교하며 기억하고 있어야 할 필요가 없게 됩니다.

또한 다 뽑고 다시 하나를 없애지 않고, 처음부터 필요한 만큼만 뽑을 수 있습니다.

 

방법2. union-find 최적화

2-1. rank를 통한 union 최적화

현재 노드를 기준으로 트리의 높이를 저장하는 rank 배열을 만듭니다. union 과정에서 트리의 높이를 기준으로 높이가 낮은 쪽으로 합치도록 구현합니다. 이런 방식으로 구현하면 트리의 높이가 한쪽으로 깊게 형성되지 않아 find 과정에서 시간을 단축시킬 수 있습니다.

public static boolean union(int a, int b) {
        a = find(a);
        b = find(b);

        if (a == b)
            return false;

        if (rank[a] < rank[b]) {
            parent[a] = b;
        } else if (rank[a] > rank[b]) {
            parent[b] = a;
        } else {
            parent[b] = a;
            rank[a]++;
        }

        return true;
}

 

2-2. 경로 압축을 통한 find 함수 최적화

find 연산을 수행할 때, 찾고자 하는 노드가 트리의 루트 노드와 연결되도록 구현했습니다. 

이런 방식으로 구현했을 경우, 같은 노드를 여러번 찾는 상황에서 처음에는 찾는 속도가 같지만, 그 이후에는 처음 찾을 때 곧바로 루트 노드에 연결되도록 구현하여 찾는 속도가 빨라집니다.

public static int find(int a) {
        if(parent[a] == a)
            return a;

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

예시

// 초기
5 -> 4 -> 3 -> 2 -> 1

// 처음 find(5) 과정
find(5) -> find(4) -> find(3) -> find(2) -> find(1) = 1

// find(5) 한번 한 후 
5, 4, 3, 2, 1 -> 5

// find(5) 다시 한 후
find(5) = 1

 

코드

import java.util.*;
import java.io.*;

/**
 * 백준 1647
 * 도시 분할 계획
 * 골드4
 * https://www.acmicpc.net/problem/1647
 */
public class Boj1647 {

    private static class Node implements Comparable<Node> {

        int from;
        int to;
        int cost;

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

        @Override
        public int compareTo(Node o) {
            return Integer.compare(cost, o.cost);
        }
    }

    static int N; // 집의 개수
    static int M; // 길의 개수
    static int[] parent; // 부모 노드
    static int[] rank;

    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());

        parent = new int[N + 1];
        rank = new int[N + 1];
        for(int i = 1; i <= N; i++) {
            parent[i] = i;
            rank[i] = 1;
        }

        PriorityQueue<Node> pq = new PriorityQueue<>();

        for(int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int to = Integer.parseInt(st.nextToken());
            int from = Integer.parseInt(st.nextToken());
            int cost = Integer.parseInt(st.nextToken());

            pq.add(new Node(to, from, cost));
        }

        int answer = 0;
        int index = 0;
        while(index < N - 2 && !pq.isEmpty()) {
            Node cur = pq.poll();

            if(union(cur.to, cur.from)) {
                answer += cur.cost;
                index++;
            }
        }

        System.out.println(answer);

    }

    public static int find(int a) {
        if(parent[a] == a)
            return a;

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

    public static boolean union(int a, int b) {
        a = find(a);
        b = find(b);

        if (a == b)
            return false;

        if (rank[a] < rank[b]) {
            parent[a] = b;
        } else if (rank[a] > rank[b]) {
            parent[b] = a;
        } else {
            parent[b] = a;
            rank[a]++;
        }

        return true;
    }
}