minimum_spanning_tree_prims.py

importsys
fromcollectionsimportdefaultdict


classHeap:
def__init__(self):
self.node_position= []

defget_position(self, vertex):
returnself.node_position[vertex]

defset_position(self, vertex, pos):
self.node_position[vertex] =pos

deftop_to_bottom(self, heap, start, size, positions):
ifstart>size//2-1:
return
else:
if2*start+2>=size: # noqa: SIM114
smallest_child=2*start+1
elifheap[2*start+1] <heap[2*start+2]:
smallest_child=2*start+1
else:
smallest_child=2*start+2
ifheap[smallest_child] <heap[start]:
temp, temp1=heap[smallest_child], positions[smallest_child]
heap[smallest_child], positions[smallest_child] = (
heap[start],
positions[start],
 )
heap[start], positions[start] =temp, temp1

temp=self.get_position(positions[smallest_child])
self.set_position(
positions[smallest_child], self.get_position(positions[start])
 )
self.set_position(positions[start], temp)

self.top_to_bottom(heap, smallest_child, size, positions)

# Update function if value of any node in min-heap decreases
defbottom_to_top(self, val, index, heap, position):
temp=position[index]

whileindex!=0:
parent=int((index-2) /2) ifindex%2==0elseint((index-1) /2)

ifval<heap[parent]:
heap[index] =heap[parent]
position[index] =position[parent]
self.set_position(position[parent], index)
else:
heap[index] =val
position[index] =temp
self.set_position(temp, index)
break
index=parent
else:
heap[0] =val
position[0] =temp
self.set_position(temp, 0)

defheapify(self, heap, positions):
start=len(heap) //2-1
foriinrange(start, -1, -1):
self.top_to_bottom(heap, i, len(heap), positions)

defdelete_minimum(self, heap, positions):
temp=positions[0]
heap[0] =sys.maxsize
self.top_to_bottom(heap, 0, len(heap), positions)
returntemp


defprisms_algorithm(adjacency_list):
"""
 >>> adjacency_list = {0: [[1, 1], [3, 3]],
 ... 1: [[0, 1], [2, 6], [3, 5], [4, 1]],
 ... 2: [[1, 6], [4, 5], [5, 2]],
 ... 3: [[0, 3], [1, 5], [4, 1]],
 ... 4: [[1, 1], [2, 5], [3, 1], [5, 4]],
 ... 5: [[2, 2], [4, 4]]}
 >>> prisms_algorithm(adjacency_list)
 [(0, 1), (1, 4), (4, 3), (4, 5), (5, 2)]
 """

heap=Heap()

visited= [0] *len(adjacency_list)
nbr_tv= [-1] *len(adjacency_list) # Neighboring Tree Vertex of selected vertex
# Minimum Distance of explored vertex with neighboring vertex of partial tree
# formed in graph
distance_tv= [] # Heap of Distance of vertices from their neighboring vertex
positions= []

forvertexinrange(len(adjacency_list)):
distance_tv.append(sys.maxsize)
positions.append(vertex)
heap.node_position.append(vertex)

tree_edges= []
visited[0] =1
distance_tv[0] =sys.maxsize
forneighbor, distanceinadjacency_list[0]:
nbr_tv[neighbor] =0
distance_tv[neighbor] =distance
heap.heapify(distance_tv, positions)

for_inrange(1, len(adjacency_list)):
vertex=heap.delete_minimum(distance_tv, positions)
ifvisited[vertex] ==0:
tree_edges.append((nbr_tv[vertex], vertex))
visited[vertex] =1
forneighbor, distanceinadjacency_list[vertex]:
if (
visited[neighbor] ==0
anddistance<distance_tv[heap.get_position(neighbor)]
 ):
distance_tv[heap.get_position(neighbor)] =distance
heap.bottom_to_top(
distance, heap.get_position(neighbor), distance_tv, positions
 )
nbr_tv[neighbor] =vertex
returntree_edges


if__name__=="__main__": # pragma: no cover
# < --------- Prims Algorithm --------- >
edges_number=int(input("Enter number of edges: ").strip())
adjacency_list=defaultdict(list)
for_inrange(edges_number):
edge= [int(x) forxininput().strip().split()]
adjacency_list[edge[0]].append([edge[1], edge[2]])
adjacency_list[edge[1]].append([edge[0], edge[2]])
print(prisms_algorithm(adjacency_list))