prim.py (1643B)
1 import heapq 2 import random 3 import math 4 5 VERTICES = 10000 6 EDGES = 100000 7 8 white = (255, 255, 255) 9 red = (255, 0, 0) 10 black = (0, 0, 0) 11 grey = (100,100,100) 12 light_grey = (50,50,50) 13 14 15 class Vertex(): 16 def __init__(self, x, y): 17 self.x = x 18 self.y = y 19 self.visited = False 20 21 class Edge(): 22 def __init__(self, v1, v2): 23 self.v1 = v1 24 self.v2 = v2 25 self.dist = math.sqrt(((v1.x - v2.x) ** 2) + ((v1.y - v2.y) ** 2)) 26 27 def __lt__(self,otr): 28 return self.dist < otr.dist 29 30 31 graph = {} 32 33 for i in range(0,VERTICES): 34 x = random.random() * 5120 35 y = random.random() * 1440 36 graph[Vertex(x,y)] = [] 37 38 39 edge_list = [] 40 for i in range(0,EDGES): 41 k1 = None 42 k2 = None 43 while k1 == k2: 44 k1 = random.choice(list(graph.keys())) 45 k2 = random.choice(list(graph.keys())) 46 47 edge = Edge(k1,k2) 48 graph[k1].append(edge) 49 graph[k2].append(edge) 50 edge_list.append(edge) 51 52 53 edge_heap = [] 54 visited_vertices = set() 55 56 mst = [] 57 58 start = random.choice(list(graph.keys())) 59 start.visited = True 60 visited_vertices.add(start) 61 for edge in graph[start]: 62 heapq.heappush(edge_heap, edge) 63 64 while True: 65 item = None 66 while edge_heap: 67 candidate = heapq.heappop(edge_heap) 68 if (candidate.v1 in visited_vertices) != (candidate.v2 in visited_vertices): 69 item = candidate 70 break 71 72 if item is None: 73 break 74 75 new_vertex = item.v2 if item.v1 in visited_vertices else item.v1 76 visited_vertices.add(new_vertex) 77 new_vertex.visited = True 78 mst.append(item) 79 for edge in graph[new_vertex]: 80 heapq.heappush(edge_heap, edge)