알고리즘

프림

쥐4 2024. 11. 13. 23:48

최소 신장 트리(MST)를 구하는 알고리즘.

 

모든 노드를 연결할때, 그 모든 연결 길이가 최소가 되게 한다.

그리디를 이용.

 

1. 시작 노드를 PriorityQueue에 저장

2. PriorityQueue가 비지 않을 때까지 && 모든 노드를 연결 할 때까지 반복문

3. PriorityQueue에서 하나 가져옴 = 최소값으로 연결할 수 있는 노드 연결

4. visited에 체크하여 이미 연결되었음을 표시

5. 한 노드를 가져왔을때, 그 노드의 진행 여부를 체크 후, 연결된 모든 노드를 priorityQueue에 추가

 

package testProject2;


import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class Main {
    static BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(System.in));
    static StringTokenizer stringTokenizer;
    static List<int[]>[] graph;
    static int testCase;
    static int nodeCount, edgeCount;

    public static void main(String args[]) throws Exception {
        testCase = Integer.parseInt(bufferedReader.readLine());
        for (int i = 0; i < testCase; i++) {
            long result = run();
            System.out.printf("#%d %d\n", i + 1, result);
        }
    }

    private static long run() throws IOException {
        initializeGraph();

        boolean[] visited = new boolean[nodeCount + 1];
        PriorityQueue<int[]> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(edge -> edge[1]));
        priorityQueue.add(new int[]{1, 0});

        long totalWeight = 0;
        int edgesUsed = 0;

        while (!priorityQueue.isEmpty() && edgesUsed < nodeCount) {
            int[] currentEdge = priorityQueue.poll();
            int currentNode = currentEdge[0];
            int weight = currentEdge[1];

            if (visited[currentNode]) continue;

            visited[currentNode] = true;
            totalWeight += weight;
            edgesUsed++;

            for (int[] neighbor : graph[currentNode]) {
                int nextNode = neighbor[0];
                int nextWeight = neighbor[1];
                if (!visited[nextNode]) {
                    priorityQueue.add(new int[]{nextNode, nextWeight});
                }
            }
        }

        if(edgesUsed == nodeCount) {
        	return totalWeight;
        }
        return -1;
    }

    private static void initializeGraph() throws IOException {
        String graphInfo = bufferedReader.readLine();
        stringTokenizer = new StringTokenizer(graphInfo);
        nodeCount = Integer.parseInt(stringTokenizer.nextToken());
        edgeCount = Integer.parseInt(stringTokenizer.nextToken());

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

        for (int i = 0; i < edgeCount; i++) {
            stringTokenizer = new StringTokenizer(bufferedReader.readLine());
            int startNode = Integer.parseInt(stringTokenizer.nextToken());
            int endNode = Integer.parseInt(stringTokenizer.nextToken());
            int value = Integer.parseInt(stringTokenizer.nextToken());

            graph[startNode].add(new int[]{endNode, value});
            graph[endNode].add(new int[]{startNode, value});
        }
    }
}