1197 - 최소 스패닝 트리

이 문제는 그래프가 주어졌을 때, 그 그래프의 최소 신장 트리(MST)를 구하는 문제이다. 최소 신장 트리란, 어떤 그래프에 대해 가능한 모든 신장 트리 중에서, 간선의 가중치 합이 가장 작은 트리를 뜻한다. 신장 트리는 그래프에서 모든 정점에 대한 최소한의 연결만을 남긴 그래프이다.

최소 신장 트리를 구하는 방법에는 여러 가지가 있는데, 필자의 경우 프림 알고리즘을 이용하여 문제를 해결하였다. 프림 알고리즘의 동작 방법은 다음과 같다.

1. 임의의 정점을 선택해 최초의 트리를 구성한다.
2. 트리에 포함된 정점에서 포함되지 않은 정점으로 이동하는 간선 중 가중치가 가장 작은 간선을 선택한다. (우선순위 큐 이용) 간선과 정점을 트리에 추가한다
3. 모든 정점이 트리에 존재할 때까지 2번을 반복한다.

🔗 문제 링크

코드

간선을 표현할 클래스 Pair을 선언한다.

static class Pair {
    int x;
    int y;

    public Pair(int x, int y) {
        this.x = x;
        this.y = y;
    }
}

그래프의 정점, 간선 개수인 V, E를 입력받는다. 그래프를 인접 리스트 형태로 저장하기 위해 graph를 선언하고 각 원소별로 ArrayList로 초기화한다.

public static void main(String[] args) throws IOException {
    StringTokenizer st = new StringTokenizer(br.readLine());
    int V, E;
    V = Integer.parseInt(st.nextToken());
    E = Integer.parseInt(st.nextToken());

    @SuppressWarnings("unchecked")
    List<Pair>[] graph = new List[V+1];

    for (int i = 1; i <= V; i++) 
        graph[i] = new ArrayList<>();

이제 그래프의 간선 정보를 입력받는다. 무향 그래프이므로 a b w 형태로 입력받은 간선 정보를 대칭적으로 저장해주어야 한다.

    for (int i = 0; i < E; i++) {
        st = new StringTokenizer(br.readLine());

        int a, b, w;
        a = Integer.parseInt(st.nextToken());
        b = Integer.parseInt(st.nextToken());
        w = Integer.parseInt(st.nextToken());

        graph[a].add(new Pair(b, w));
        graph[b].add(new Pair(a, w));
    }

프림 알고리즘 수행을 위해 필요한 변수들을 선언한다. visited는 정점이 최소 신장 트리에 포함되었는지 확인하기 위한 변수이다. minWeight는 최소 신장 트리의 간선 가중치의 총합을 저장하는 변수이고, selected는 최소 신장 트리에 포함된 정점의 개수를 세기 위한 변수이다. selected=-1로 초기화한 이유는 제일 처음 넣은 간선이 실제로 존재하는 것이 아니고 1을 최초 정점으로 두어 처리하기 위한 임의의 간선이기 때문에, 선택된 간선 개수에서 제외하기 위한 것이다.

최소 신장 트리에 포함된 정점으로부터 접근 가능한 간선 중, 가중치가 가장 작은 간선을 $O(\log V)$ 에 확인하기 위한 우선순위 큐를 선언한다. 우선순위 큐는 y(가중치)를 기준으로 최소 힙 구조를 가지도록 설정한다. 또한 최초로 방문할 간선을 임의로 삽입한다.

    boolean[] visited = new boolean[V+1];
    int minWeight = 0; 
    int selected = -1;
    PriorityQueue<Pair> pq = new PriorityQueue<>((a, b)->Long.compare(a.y, b.y));
    pq.offer(new Pair(1, 0));

이제 반복문 부분을 보자. 내용이 다익스트라 알고리즘과 유사하다는 것을 알 수 있다. 먼저 우선순위 큐에서 후보 간선을 하나 꺼낸 뒤, 간선이 이미 확정된 정점으로의 이동이라면 건너뛴다. 확정되지 않은 정점에 대한 이동이면, 정점을 확정하고 minWeight에 가중치를 더한다. 그런 다음 이 간선을 통해 새로 확정한 정점에 대해, 그 정점으로부터 인접한 간선을 탐색하며 우선순위 큐에 추가한다. 이 때 만약 확정된 정점으로 이동시키는 간선이라면 추가하지 않는다.

만약 선택된 간선의 개수가 V-1과 같아지면 모든 정점과 그것을 트리가 되도록 연결하는 간선을 선택한 것이므로 반복문을 빠져나오도록 한다.

참고로 이 부분을 구현할 때 조심해서 생각해야 할 부분은, 우선순위 큐에 간선을 넣는다고 해서 그 간선이 확정이 된 것은 아니라는 점이다. 다익스트라 알고리즘을 알고 있다면 어떤 뜻인지 이해가 될 것이다. 우선순위 큐에 넣은 것은 후보 간선이고, 후보들 중 최소 가중치를 가진 간선을 우선순위 큐에서 꺼낸 그 순간에 확정되는 것이므로, 우선순위 큐에서 꺼낸 후에 간선을 확정지어야 한다.

    while (!pq.isEmpty()) {
        Pair s = pq.poll();
        if (visited[s.x]) continue;
        visited[s.x] = true;
        minWeight += s.y;
        if ((++selected) == V-1) break;

        for (Pair adj : graph[s.x]) {
            if (visited[adj.x]) continue;
            pq.offer(adj);
        }
    }

반복문을 빠져나왔다면 최소 신장 트리를 완성한 것이므로 가중치를 출력하고 프로그램을 종료한다.

    System.out.println(minWeight);
}

전체 코드

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

class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

    static class Pair {
        int x;
        int y;

        public Pair(int x, int y) {
            this.x = x;
            this.y = y;
        }
    }
    public static void main(String[] args) throws IOException {
        StringTokenizer st = new StringTokenizer(br.readLine());
        int V, E;
        V = Integer.parseInt(st.nextToken());
        E = Integer.parseInt(st.nextToken());

        @SuppressWarnings("unchecked")
        List<Pair>[] graph = new List[V+1];

        for (int i = 1; i <= V; i++) 
            graph[i] = new ArrayList<>();

        for (int i = 0; i < E; i++) {
            st = new StringTokenizer(br.readLine());

            int a, b, w;
            a = Integer.parseInt(st.nextToken());
            b = Integer.parseInt(st.nextToken());
            w = Integer.parseInt(st.nextToken());

            graph[a].add(new Pair(b, w));
            graph[b].add(new Pair(a, w));
        }

        boolean[] visited = new boolean[V+1];
        int minWeight = 0; 
        int selected = -1;
        PriorityQueue<Pair> pq = new PriorityQueue<>((a, b)->Long.compare(a.y, b.y));
        pq.offer(new Pair(1, 0));

        while (!pq.isEmpty()) {
            Pair s = pq.poll();
            if (visited[s.x]) continue;
            visited[s.x] = true;
            minWeight += s.y;
            if ((++selected) == V-1) break;

            for (Pair adj : graph[s.x]) {
                if (visited[adj.x]) continue;
                pq.offer(adj);
            }
        }
        System.out.println(minWeight);
    }
}

시간복잡도

while 반복문 자체는 모든 간선을 확인하기 때문에 약 $E$번 순회할 것이다. 또한 우선순위 큐에 들어가는 간선도 $E$개일 것이다. 따라서 전체 시간복잡도는 $O(E \log E)$이다. (우선순위 큐의 삽입/삭제 시간복잡도는 $O(\log E)$)