visualizations

Programmatic visualizations
git clone git://git.laack.co/visualizations.git
Log | Files | Refs | README

prim.py (1914B)

      1 import heapq
      2 import random
      3 import math
      4 
      5 VERTICES = 200_000
      6 EDGES = 2_000_000
      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 class Edge():
     21     def __init__(self, v1, v2):
     22         self.v1 = v1
     23         self.v2 = v2
     24         self.dist = math.sqrt(((v1.x - v2.x) ** 2) + ((v1.y - v2.y) ** 2))
     25     def __lt__(self,otr):
     26         return self.dist < otr.dist
     27 
     28 
     29 graph = {}
     30 
     31 for i in range(0,VERTICES):
     32     x = random.random() * 5120
     33     y = random.random() * 1440
     34     graph[Vertex(x,y)] = []
     35 
     36 
     37 edge_list = []
     38 keys = list(graph.keys())
     39 
     40 for i in range(0,EDGES):
     41     k1 = None
     42     k2 = None
     43     while k1 == k2:
     44         k1 = random.choice(keys)
     45         k2 = random.choice(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 
     57 start = random.choice(keys)
     58 start.visited = True
     59 visited_vertices.add(start)
     60 for edge in graph[start]:
     61     heapq.heappush(edge_heap, edge)
     62 
     63 
     64 first = True
     65 
     66 # we still have this for consistency with the c++ variant because the c++ variant is also tracking this info.
     67 to_draw_vert = []
     68 to_draw_edge = []
     69 
     70 while True:
     71     item = None
     72     while edge_heap:
     73         candidate = heapq.heappop(edge_heap)
     74         if (candidate.v1 in visited_vertices) != (candidate.v2 in visited_vertices):
     75             item = candidate
     76             break
     77 
     78     if item is None:
     79         break
     80 
     81     new_vertex = item.v2 if item.v1 in visited_vertices else item.v1
     82     visited_vertices.add(new_vertex)
     83     new_vertex.visited = True
     84 
     85     to_draw_vert.append(new_vertex)
     86     to_draw_edge.append(item)
     87 
     88     for edge in graph[new_vertex]:
     89         if not edge.v1 in visited_vertices or not edge.v2 in visited_vertices:
     90             heapq.heappush(edge_heap, edge)