Prim's (MST) : Special Subtree

  • + 0 comments

    Java, O(log E)

    public static int prims(int n, List<List<Integer>> edges, int start) {
                Map<Integer, List<int[]>> graph = new HashMap<>();
                for (List<Integer> edge : edges) {
                    graph.computeIfAbsent(edge.get(0), k -> new ArrayList<>()).add(new int[]{edge.get(1), edge.get(2)});
                    graph.computeIfAbsent(edge.get(1), k -> new ArrayList<>()).add(new int[]{edge.get(0), edge.get(2)});
                }
    
                PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(a -> a[1]));
                pq.add(new int[]{start, 0}); 
    
                boolean[] visited = new boolean[n + 1];
                int totalWeight = 0;
    
                while (!pq.isEmpty()) {
                    int[] current = pq.poll();
                    int currentNode = current[0];
                    int edgeWeight = current[1];
    
                    if (visited[currentNode]) {
                        continue;
                    }
    
                    visited[currentNode] = true;
                    totalWeight += edgeWeight;
    
                    if (graph.containsKey(currentNode)) {
                        for (int[] neighbor : graph.get(currentNode)) {
                            if (!visited[neighbor[0]]) {
                                pq.add(new int[]{neighbor[0], neighbor[1]});
                            }
                        }
                    }
                }
    
                return totalWeight;
            }