Prim's algorithm grows the MST from a single starting node instead of sorting all edges. It repeatedly pulls the cheapest edge that reaches a not-yet-visited node, using a min-heap to always pop the smallest candidate. For the same graph it yields the same total weight as Kruskal's.
Implement prim(n, adj):
Real use: Prim's is often preferred for dense networks where an adjacency structure is already on hand.
import heapq
class DSU:
def __init__(self, n):
self.parent = list(range(n))
self.rank = [0] * n
def find(self, x):
while self.parent[x] != x:
self.parent[x] = self.parent[self.parent[x]]
x = self.parent[x]
return x
def union(self, a, b):
ra, rb = self.find(a), self.find(b)
if ra == rb:
return False
if self.rank[ra] < self.rank[rb]:
ra, rb = rb, ra
self.parent[rb] = ra
if self.rank[ra] == self.rank[rb]:
self.rank[ra] += 1
return True
def connected(self, a, b):
return self.find(a) == self.find(b)
def kruskal(n, edges):
dsu = DSU(n)
total = 0
for w, u, v in sorted((w, u, v) for (u, v, w) in edges):
if dsu.union(u, v):
total += w
return total
def kruskal_edges(n, edges):
dsu = DSU(n)
chosen = []
for w, u, v in sorted((w, u, v) for (u, v, w) in edges):
if dsu.union(u, v):
chosen.append((u, v, w))
return chosen
def prim(n, adj):
"""Return the total weight of the MST using Prim's algorithm and a heap."""
# TODO: start at node 0, pop the cheapest edge to an unvisited node, sum weights
pass
The editor and the test runner need a wide screen. Open this page on a desktop browser and your progress will be waiting.
Prim's algorithm grows the MST from a single starting node instead of sorting all edges. It repeatedly pulls the cheapest edge that reaches a not-yet-visited node, using a min-heap to always pop the smallest candidate. For the same graph it yields the same total weight as Kruskal's.
Implement prim(n, adj):
Real use: Prim's is often preferred for dense networks where an adjacency structure is already on hand.
Press Run Tests or ⌘↩ to check your solution.