minimum_spanning_tree_prims2.py

"""
Prim's (also known as Jarník's) algorithm is a greedy algorithm that finds a minimum
spanning tree for a weighted undirected graph. This means it finds a subset of the
edges that forms a tree that includes every vertex, where the total weight of all the
edges in the tree is minimized. The algorithm operates by building this tree one vertex
at a time, from an arbitrary starting vertex, at each step adding the cheapest possible
connection from the tree to another vertex.
"""

from __future__ importannotations

fromsysimportmaxsize
fromtypingimportGeneric, TypeVar

T=TypeVar("T")


defget_parent_position(position: int) ->int:
"""
 heap helper function get the position of the parent of the current node

 >>> get_parent_position(1)
 0
 >>> get_parent_position(2)
 0
 """
return (position-1) //2


defget_child_left_position(position: int) ->int:
"""
 heap helper function get the position of the left child of the current node

 >>> get_child_left_position(0)
 1
 """
return (2*position) +1


defget_child_right_position(position: int) ->int:
"""
 heap helper function get the position of the right child of the current node

 >>> get_child_right_position(0)
 2
 """
return (2*position) +2


classMinPriorityQueue(Generic[T]):
"""
 Minimum Priority Queue Class

 Functions:
 is_empty: function to check if the priority queue is empty
 push: function to add an element with given priority to the queue
 extract_min: function to remove and return the element with lowest weight (highest
 priority)
 update_key: function to update the weight of the given key
 _bubble_up: helper function to place a node at the proper position (upward
 movement)
 _bubble_down: helper function to place a node at the proper position (downward
 movement)
 _swap_nodes: helper function to swap the nodes at the given positions

 >>> queue = MinPriorityQueue()

 >>> queue.push(1, 1000)
 >>> queue.push(2, 100)
 >>> queue.push(3, 4000)
 >>> queue.push(4, 3000)

 >>> queue.extract_min()
 2

 >>> queue.update_key(4, 50)

 >>> queue.extract_min()
 4
 >>> queue.extract_min()
 1
 >>> queue.extract_min()
 3
 """

def__init__(self) ->None:
self.heap: list[tuple[T, int]] = []
self.position_map: dict[T, int] = {}
self.elements: int=0

def__len__(self) ->int:
returnself.elements

def__repr__(self) ->str:
returnstr(self.heap)

defis_empty(self) ->bool:
# Check if the priority queue is empty
returnself.elements==0

defpush(self, elem: T, weight: int) ->None:
# Add an element with given priority to the queue
self.heap.append((elem, weight))
self.position_map[elem] =self.elements
self.elements+=1
self._bubble_up(elem)

defextract_min(self) ->T:
# Remove and return the element with lowest weight (highest priority)
ifself.elements>1:
self._swap_nodes(0, self.elements-1)
elem, _=self.heap.pop()
delself.position_map[elem]
self.elements-=1
ifself.elements>0:
bubble_down_elem, _=self.heap[0]
self._bubble_down(bubble_down_elem)
returnelem

defupdate_key(self, elem: T, weight: int) ->None:
# Update the weight of the given key
position=self.position_map[elem]
self.heap[position] = (elem, weight)
ifposition>0:
parent_position=get_parent_position(position)
_, parent_weight=self.heap[parent_position]
ifparent_weight>weight:
self._bubble_up(elem)
else:
self._bubble_down(elem)
else:
self._bubble_down(elem)

def_bubble_up(self, elem: T) ->None:
# Place a node at the proper position (upward movement) [to be used internally
# only]
curr_pos=self.position_map[elem]
ifcurr_pos==0:
returnNone
parent_position=get_parent_position(curr_pos)
_, weight=self.heap[curr_pos]
_, parent_weight=self.heap[parent_position]
ifparent_weight>weight:
self._swap_nodes(parent_position, curr_pos)
returnself._bubble_up(elem)
returnNone

def_bubble_down(self, elem: T) ->None:
# Place a node at the proper position (downward movement) [to be used
# internally only]
curr_pos=self.position_map[elem]
_, weight=self.heap[curr_pos]
child_left_position=get_child_left_position(curr_pos)
child_right_position=get_child_right_position(curr_pos)
ifchild_left_position<self.elementsandchild_right_position<self.elements:
_, child_left_weight=self.heap[child_left_position]
_, child_right_weight=self.heap[child_right_position]
ifchild_right_weight<child_left_weightandchild_right_weight<weight:
self._swap_nodes(child_right_position, curr_pos)
returnself._bubble_down(elem)
ifchild_left_position<self.elements:
_, child_left_weight=self.heap[child_left_position]
ifchild_left_weight<weight:
self._swap_nodes(child_left_position, curr_pos)
returnself._bubble_down(elem)
else:
returnNone
ifchild_right_position<self.elements:
_, child_right_weight=self.heap[child_right_position]
ifchild_right_weight<weight:
self._swap_nodes(child_right_position, curr_pos)
returnself._bubble_down(elem)
returnNone

def_swap_nodes(self, node1_pos: int, node2_pos: int) ->None:
# Swap the nodes at the given positions
node1_elem=self.heap[node1_pos][0]
node2_elem=self.heap[node2_pos][0]
self.heap[node1_pos], self.heap[node2_pos] = (
self.heap[node2_pos],
self.heap[node1_pos],
 )
self.position_map[node1_elem] =node2_pos
self.position_map[node2_elem] =node1_pos


classGraphUndirectedWeighted(Generic[T]):
"""
 Graph Undirected Weighted Class

 Functions:
 add_node: function to add a node in the graph
 add_edge: function to add an edge between 2 nodes in the graph
 """

def__init__(self) ->None:
self.connections: dict[T, dict[T, int]] = {}
self.nodes: int=0

def__repr__(self) ->str:
returnstr(self.connections)

def__len__(self) ->int:
returnself.nodes

defadd_node(self, node: T) ->None:
# Add a node in the graph if it is not in the graph
ifnodenotinself.connections:
self.connections[node] = {}
self.nodes+=1

defadd_edge(self, node1: T, node2: T, weight: int) ->None:
# Add an edge between 2 nodes in the graph
self.add_node(node1)
self.add_node(node2)
self.connections[node1][node2] =weight
self.connections[node2][node1] =weight


defprims_algo(
graph: GraphUndirectedWeighted[T],
) ->tuple[dict[T, int], dict[T, T|None]]:
"""
 >>> graph = GraphUndirectedWeighted()

 >>> graph.add_edge("a", "b", 3)
 >>> graph.add_edge("b", "c", 10)
 >>> graph.add_edge("c", "d", 5)
 >>> graph.add_edge("a", "c", 15)
 >>> graph.add_edge("b", "d", 100)

 >>> dist, parent = prims_algo(graph)

 >>> abs(dist["a"] - dist["b"])
 3
 >>> abs(dist["d"] - dist["b"])
 15
 >>> abs(dist["a"] - dist["c"])
 13
 """
# prim's algorithm for minimum spanning tree
dist: dict[T, int] =dict.fromkeys(graph.connections, maxsize)
parent: dict[T, T|None] =dict.fromkeys(graph.connections)

priority_queue: MinPriorityQueue[T] =MinPriorityQueue()
fornode, weightindist.items():
priority_queue.push(node, weight)

ifpriority_queue.is_empty():
returndist, parent

# initialization
node=priority_queue.extract_min()
dist[node] =0
forneighbouringraph.connections[node]:
ifdist[neighbour] >dist[node] +graph.connections[node][neighbour]:
dist[neighbour] =dist[node] +graph.connections[node][neighbour]
priority_queue.update_key(neighbour, dist[neighbour])
parent[neighbour] =node

# running prim's algorithm
whilenotpriority_queue.is_empty():
node=priority_queue.extract_min()
forneighbouringraph.connections[node]:
ifdist[neighbour] >dist[node] +graph.connections[node][neighbour]:
dist[neighbour] =dist[node] +graph.connections[node][neighbour]
priority_queue.update_key(neighbour, dist[neighbour])
parent[neighbour] =node
returndist, parent