Implement Prims Algorithm

The problem, simply

You get V towns and possible two-way roads, each with a building price. Pick roads so every town can reach every other town, for the smallest total price. Give back that total.

This is the same goal as Kruskal's algorithm, the cheapest connected network. Prim reaches it differently: instead of looking at all roads at once, it grows one connected blob from a starting town.

A story to remember it

Picture it

Examples

Way 1: the simple way

Keep a price tag on every town outside the blot: the cheapest road from the blot to that town found so far (its key). Each turn, scan all towns to find the dry one with the smallest tag.

1. Every tag starts at infinity, except town 0 at 0.
2. Repeat V times: scan all towns for the dry one with the smallest tag. Ink it and add its tag to the total.
3. For each road from the new town to a dry town, lower that town's tag if the road is cheaper.

import java.util.*;

class Solution {
    public int spanningTree(int V, int[][] edges) {
        List<List<int[]>> adj = new ArrayList<>();
        for (int i = 0; i < V; i++) adj.add(new ArrayList<>());
        for (int[] e : edges) {
            adj.get(e[0]).add(new int[]{e[1], e[2]});
            adj.get(e[1]).add(new int[]{e[0], e[2]});
        }
        int[] key = new int[V];                    // cheapest road from the blot
        Arrays.fill(key, Integer.MAX_VALUE);
        key[0] = 0;
        boolean[] inked = new boolean[V];
        int cost = 0;
        for (int turn = 0; turn < V; turn++) {
            int u = -1;
            for (int v = 0; v < V; v++) {          // scan every town
                if (!inked[v] && (u == -1 || key[v] < key[u])) u = v;
            }
            inked[u] = true;
            cost += key[u];
            for (int[] nb : adj.get(u)) {
                if (!inked[nb[0]] && nb[1] < key[nb[0]]) key[nb[0]] = nb[1];
            }
        }
        return cost;
    }
}

How long it takes: V turns, each scanning V towns, O(V^2) (the roads add O(E), which is at most V^2). Extra room: the tags and flags, O(V). This is actually the best choice when almost every pair of towns has a road.

Way 2: the clever way

Put the offers in a min-heap, a bag that always gives back the cheapest offer. An offer is "town v can join for price w".

1. Push the offer (town 0, price 0).
2. Pop the cheapest offer. If that town is already inked, the offer is old: throw it away.
3. Otherwise ink the town and add the price. Push an offer for each road from it to a dry town.
4. Stop when all V towns are inked.

We do not delete old offers when a cheaper one shows up. We let them sit in the heap and skip them when they come out. This is called lazy deletion, and it keeps the code short.

import java.util.*;

class Solution {
    public int spanningTree(int V, int[][] edges) {
        List<List<int[]>> adj = new ArrayList<>();
        for (int i = 0; i < V; i++) adj.add(new ArrayList<>());
        for (int[] e : edges) {
            adj.get(e[0]).add(new int[]{e[1], e[2]});
            adj.get(e[1]).add(new int[]{e[0], e[2]});
        }
        boolean[] inked = new boolean[V];
        PriorityQueue<int[]> heap = new PriorityQueue<>((x, y) -> Integer.compare(x[0], y[0]));  // {price, town}
        heap.add(new int[]{0, 0});
        int cost = 0, joined = 0;
        while (!heap.isEmpty() && joined < V) {
            int[] top = heap.poll();
            int w = top[0], u = top[1];
            if (inked[u]) continue;               // an old offer, skip it
            inked[u] = true;
            cost += w;
            joined++;
            for (int[] nb : adj.get(u)) {
                if (!inked[nb[0]]) heap.add(new int[]{nb[1], nb[0]});
            }
        }
        return cost;
    }
}

How long it takes: every road can add one offer from each end, so the heap sees O(E) pushes and pops, each O(log E). Since log E is at most 2 log V, that is O(E log V). Extra room: the neighbour lists and the heap, O(V + E).

Watch out for

• Treating it like Dijkstra. Dijkstra pushes the total distance from the start (dist[u] + w). Prim pushes only the price of the one road, w.
• A map in pieces. The blot never reaches the other piece, so joined stays below V. Check it if you must report that.
• Building the neighbour list one way only. Roads are two-way, so add both directions.
• Stopping when the heap is empty is fine too, but stopping at joined == V saves popping old offers.

What interviewers ask next

• Prim or Kruskal? On a sparse map both are about O(E log V). On a dense map, the simple O(V^2) Prim from Way 1 is the fastest, with no heap at all.
• Can you print the roads, not only the total? Push the town the offer came from too, {price, town, from}, and record from-town when a town joins.
• Why is it correct? The cheapest road leaving the blot crosses the split "blot versus the rest", and the cheapest road across any split is safe to take (the cut property).
• Minimum cost to connect points on a plane (LeetCode 1584)? Every pair of points is a road, so the map is dense: use the O(V^2) version.
• Can you do better than O(E log V)? A Fibonacci heap with decrease-key gives O(E + V log V), but it is rarely used in practice.

Remember this

Grown-up words: towns are vertices, roads are weighted edges, the blot is the growing tree, the magic bag is a min-heap (priority queue), the price tags are keys, skipping stale offers is lazy deletion, the result is a minimum spanning tree (MST), and the method is Prim's algorithm.