Python/graphs/minimum_spanning_tree_prims.py

116 lines
3.7 KiB
Python
Raw Normal View History

2018-10-19 12:48:28 +00:00
import sys
from collections import defaultdict
2019-10-05 05:14:13 +00:00
def PrimsAlgorithm(l): # noqa: E741
2018-10-19 12:48:28 +00:00
nodePosition = []
2019-10-05 05:14:13 +00:00
def get_position(vertex):
2018-10-19 12:48:28 +00:00
return nodePosition[vertex]
def set_position(vertex, pos):
2018-10-19 12:48:28 +00:00
nodePosition[vertex] = pos
def top_to_bottom(heap, start, size, positions):
2018-10-19 12:48:28 +00:00
if start > size // 2 - 1:
return
else:
if 2 * start + 2 >= size:
m = 2 * start + 1
else:
if heap[2 * start + 1] < heap[2 * start + 2]:
m = 2 * start + 1
else:
m = 2 * start + 2
if heap[m] < heap[start]:
temp, temp1 = heap[m], positions[m]
heap[m], positions[m] = heap[start], positions[start]
heap[start], positions[start] = temp, temp1
temp = get_position(positions[m])
set_position(positions[m], get_position(positions[start]))
set_position(positions[start], temp)
2018-10-19 12:48:28 +00:00
top_to_bottom(heap, m, size, positions)
2018-10-19 12:48:28 +00:00
# Update function if value of any node in min-heap decreases
def bottom_to_top(val, index, heap, position):
2018-10-19 12:48:28 +00:00
temp = position[index]
2019-10-05 05:14:13 +00:00
while index != 0:
2018-10-19 12:48:28 +00:00
if index % 2 == 0:
2019-10-05 05:14:13 +00:00
parent = int((index - 2) / 2)
2018-10-19 12:48:28 +00:00
else:
2019-10-05 05:14:13 +00:00
parent = int((index - 1) / 2)
2018-10-19 12:48:28 +00:00
if val < heap[parent]:
heap[index] = heap[parent]
position[index] = position[parent]
set_position(position[parent], index)
2018-10-19 12:48:28 +00:00
else:
heap[index] = val
position[index] = temp
set_position(temp, index)
2018-10-19 12:48:28 +00:00
break
index = parent
else:
heap[0] = val
position[0] = temp
set_position(temp, 0)
2018-10-19 12:48:28 +00:00
def heapify(heap, positions):
start = len(heap) // 2 - 1
for i in range(start, -1, -1):
top_to_bottom(heap, i, len(heap), positions)
2018-10-19 12:48:28 +00:00
def deleteMinimum(heap, positions):
temp = positions[0]
heap[0] = sys.maxsize
top_to_bottom(heap, 0, len(heap), positions)
2018-10-19 12:48:28 +00:00
return temp
visited = [0 for i in range(len(l))]
2019-10-05 05:14:13 +00:00
Nbr_TV = [-1 for i in range(len(l))] # Neighboring Tree Vertex of selected vertex
2018-10-19 12:48:28 +00:00
# Minimum Distance of explored vertex with neighboring vertex of partial tree formed in graph
2019-10-05 05:14:13 +00:00
Distance_TV = [] # Heap of Distance of vertices from their neighboring vertex
2018-10-19 12:48:28 +00:00
Positions = []
for x in range(len(l)):
p = sys.maxsize
Distance_TV.append(p)
Positions.append(x)
nodePosition.append(x)
TreeEdges = []
visited[0] = 1
Distance_TV[0] = sys.maxsize
for x in l[0]:
2019-10-05 05:14:13 +00:00
Nbr_TV[x[0]] = 0
Distance_TV[x[0]] = x[1]
2018-10-19 12:48:28 +00:00
heapify(Distance_TV, Positions)
for i in range(1, len(l)):
vertex = deleteMinimum(Distance_TV, Positions)
if visited[vertex] == 0:
TreeEdges.append((Nbr_TV[vertex], vertex))
visited[vertex] = 1
for v in l[vertex]:
if visited[v[0]] == 0 and v[1] < Distance_TV[get_position(v[0])]:
Distance_TV[get_position(v[0])] = v[1]
bottom_to_top(v[1], get_position(v[0]), Distance_TV, Positions)
2019-10-05 05:14:13 +00:00
Nbr_TV[v[0]] = vertex
2018-10-19 12:48:28 +00:00
return TreeEdges
2019-10-05 05:14:13 +00:00
if __name__ == "__main__":
# < --------- Prims Algorithm --------- >
n = int(input("Enter number of vertices: ").strip())
e = int(input("Enter number of edges: ").strip())
adjlist = defaultdict(list)
for x in range(e):
l = [int(x) for x in input().strip().split()] # noqa: E741
2019-10-05 05:14:13 +00:00
adjlist[l[0]].append([l[1], l[2]])
adjlist[l[1]].append([l[0], l[2]])
print(PrimsAlgorithm(adjlist))