prim.py (1727B)
1 import heapq 2 import random 3 import math 4 5 VERTICES = 100000 6 EDGES = 1000000 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 keys = list(graph.keys()) 41 42 for i in range(0,EDGES): 43 k1 = None 44 k2 = None 45 while k1 == k2: 46 k1 = random.choice(keys) 47 k2 = random.choice(keys) 48 49 edge = Edge(k1,k2) 50 graph[k1].append(edge) 51 graph[k2].append(edge) 52 edge_list.append(edge) 53 54 55 edge_heap = [] 56 visited_vertices = set() 57 58 mst = [] 59 60 start = random.choice(list(graph.keys())) 61 start.visited = True 62 visited_vertices.add(start) 63 for edge in graph[start]: 64 heapq.heappush(edge_heap, edge) 65 66 while True: 67 item = None 68 while edge_heap: 69 candidate = heapq.heappop(edge_heap) 70 if (candidate.v1 in visited_vertices) != (candidate.v2 in visited_vertices): 71 item = candidate 72 break 73 74 if item is None: 75 break 76 77 new_vertex = item.v2 if item.v1 in visited_vertices else item.v1 78 visited_vertices.add(new_vertex) 79 new_vertex.visited = True 80 mst.append(item) 81 for edge in graph[new_vertex]: 82 if not edge.v1 in visited_vertices or not edge.v2 in visited_vertices: 83 heapq.heappush(edge_heap, edge)