/
prims.py
55 lines (35 loc) · 1.05 KB
/
prims.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import math
import heapq
def printMST(parent, graph, v):
print("Edge \tWeight")
for i in range(1, v):
print(parent[i], "-", i, "\t", graph[i][ parent[i] ])
def prims(adjMatrix, v): #adjMatrix is a 2d list
inMST = [False] * v
priority = [math.inf] * v
predecessor = [None] * v
priority[0] = 0
predecessor[0] = -1
for i in range(v):
min = math.inf
minIndex = 0
for j in range(v):
if priority[j] < min and inMST[j] == False:
min = priority[j]
minIndex = j
inMST[minIndex] = True
for j in range(v):
if adjMatrix[minIndex][j] > 0 and inMST[j] == False and priority[j] > adjMatrix[minIndex][j]:
priority[j] = adjMatrix[minIndex][j]
predecessor[j] = minIndex
printMST(predecessor, adjMatrix, v)
# def heapPrims(adjList, v):
# inMST = [False] * v
# predecessor = [None] * v
# priority = [math.inf] * v
# priority[0] = 0
# heapq.heapify(priority)
# while len(priority) > 0:
# min = heapq.heappop(priority)
matrix = [[0, 5, 1, 3],[5, 0, 3, 1],[1, 3, 0, 2],[3, 1, 2, 0]]
prims(matrix, 4)