카테고리 없음

Prim's(MST): Special Subtree

hyunmin43240 2025. 5. 5. 15:59

https://www.hackerrank.com/challenges/primsmstsub/problem?isFullScreen=true

public static int prims(int n, List<List<Integer>> edges, int start) {
    // Write your code here
        
        List<List<int[]>> graph = new ArrayList<>();
    for (int i = 0; i <= n; i++) graph.add(new ArrayList<>());
    for (List<Integer> edge : edges) {
        int u = edge.get(0), v = edge.get(1), w = edge.get(2);
        graph.get(u).add(new int[]{v, w});
        graph.get(v).add(new int[]{u, w}); 
    }

    boolean[] visited = new boolean[n + 1];
    PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(a -> a[1]));
    pq.offer(new int[]{start, 0});
    int totalWeight = 0;

    while (!pq.isEmpty()) {
        int[] curr = pq.poll();
        int node = curr[0], weight = curr[1];
        if (visited[node]) continue;
        visited[node] = true;
        totalWeight += weight;

        for (int[] neighbor : graph.get(node)) {
            int nextNode = neighbor[0], nextWeight = neighbor[1];
            if (!visited[nextNode]) {
                pq.offer(new int[]{nextNode, nextWeight});
            }
        }
    }

    return totalWeight;
        
    }