Python/graphs/prim.py

153 lines
3.4 KiB
Python
Raw Normal View History

"""Prim's Algorithm.
Determines the minimum spanning tree(MST) of a graph using the Prim's Algorithm.
2019-05-26 16:41:46 +00:00
Details: https://en.wikipedia.org/wiki/Prim%27s_algorithm
2019-05-26 16:41:46 +00:00
"""
import heapq as hq
2019-05-26 16:41:46 +00:00
import math
from collections.abc import Iterator
2019-05-26 16:41:46 +00:00
class Vertex:
2019-05-26 16:41:46 +00:00
"""Class Vertex."""
def __init__(self, id):
"""
Arguments:
id - input an id to identify the vertex
Attributes:
neighbors - a list of the vertices it is linked to
edges - a dict to store the edges's weight
"""
self.id = str(id)
self.key = None
self.pi = None
self.neighbors = []
self.edges = {} # {vertex:distance}
2019-05-26 16:41:46 +00:00
def __lt__(self, other):
"""Comparison rule to < operator."""
2019-10-05 05:14:13 +00:00
return self.key < other.key
2019-05-26 16:41:46 +00:00
def __repr__(self):
"""Return the vertex id."""
return self.id
def add_neighbor(self, vertex):
2019-05-26 16:41:46 +00:00
"""Add a pointer to a vertex at neighbor's list."""
self.neighbors.append(vertex)
def add_edge(self, vertex, weight):
2019-05-26 16:41:46 +00:00
"""Destination vertex and weight."""
self.edges[vertex.id] = weight
def connect(graph, a, b, edge):
# add the neighbors:
graph[a - 1].add_neighbor(graph[b - 1])
graph[b - 1].add_neighbor(graph[a - 1])
# add the edges:
graph[a - 1].add_edge(graph[b - 1], edge)
graph[b - 1].add_edge(graph[a - 1], edge)
def prim(graph: list, root: Vertex) -> list:
"""Prim's Algorithm.
Runtime:
O(mn) with `m` edges and `n` vertices
Return:
List with the edges of a Minimum Spanning Tree
Usage:
prim(graph, graph[0])
2019-05-26 16:41:46 +00:00
"""
a = []
2019-05-26 16:41:46 +00:00
for u in graph:
u.key = math.inf
u.pi = None
root.key = 0
q = graph[:]
while q:
u = min(q)
q.remove(u)
2019-05-26 16:41:46 +00:00
for v in u.neighbors:
if (v in q) and (u.edges[v.id] < v.key):
2019-05-26 16:41:46 +00:00
v.pi = u
v.key = u.edges[v.id]
for i in range(1, len(graph)):
a.append((int(graph[i].id) + 1, int(graph[i].pi.id) + 1))
return a
def prim_heap(graph: list, root: Vertex) -> Iterator[tuple]:
"""Prim's Algorithm with min heap.
Runtime:
O((m + n)log n) with `m` edges and `n` vertices
Yield:
Edges of a Minimum Spanning Tree
Usage:
prim(graph, graph[0])
"""
for u in graph:
u.key = math.inf
u.pi = None
root.key = 0
h = list(graph)
hq.heapify(h)
while h:
u = hq.heappop(h)
for v in u.neighbors:
if (v in h) and (u.edges[v.id] < v.key):
v.pi = u
v.key = u.edges[v.id]
hq.heapify(h)
for i in range(1, len(graph)):
yield (int(graph[i].id) + 1, int(graph[i].pi.id) + 1)
def test_vector() -> None:
"""
# Creates a list to store x vertices.
>>> x = 5
>>> G = [Vertex(n) for n in range(x)]
>>> connect(G, 1, 2, 15)
>>> connect(G, 1, 3, 12)
>>> connect(G, 2, 4, 13)
>>> connect(G, 2, 5, 5)
>>> connect(G, 3, 2, 6)
>>> connect(G, 3, 4, 6)
>>> connect(G, 0, 0, 0) # Generate the minimum spanning tree:
>>> G_heap = G[:]
>>> MST = prim(G, G[0])
>>> MST_heap = prim_heap(G, G[0])
>>> for i in MST:
... print(i)
(2, 3)
(3, 1)
(4, 3)
(5, 2)
>>> for i in MST_heap:
... print(i)
(2, 3)
(3, 1)
(4, 3)
(5, 2)
"""
if __name__ == "__main__":
import doctest
doctest.testmod()