visualizations

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

prim.py (2824B)

      1 import pygame
      2 import heapq
      3 import random
      4 import math
      5 
      6 pygame.init()
      7 
      8 display = pygame.display.set_mode((1920,1080))
      9 
     10 VERTICES = 250
     11 EDGES = 1000
     12 
     13 white = (255, 255, 255)
     14 red = (255, 0, 0)
     15 black = (0, 0, 0)
     16 grey = (100,100,100)
     17 light_grey = (50,50,50)
     18 
     19 
     20 class Vertex():
     21     def __init__(self, x, y):
     22         self.x = x
     23         self.y = y
     24         self.visited = False
     25     def draw_vertex(self,display):
     26         if self.visited:
     27             pygame.draw.circle(display, white, (self.x,self.y), 5)
     28         else:
     29             pygame.draw.circle(display, grey, (self.x,self.y), 5)
     30 class Edge():
     31     def __init__(self, v1, v2):
     32         self.v1 = v1
     33         self.v2 = v2
     34         self.dist = math.sqrt(((v1.x - v2.x) ** 2) + ((v1.y - v2.y) ** 2))
     35     def draw_edge(self,display, c):
     36         pygame.draw.line(display, c, (self.v1.x,self.v1.y), (self.v2.x,self.v2.y), 1)
     37 
     38     def __lt__(self,otr):
     39         return self.dist < otr.dist
     40 
     41 
     42 graph = {}
     43 
     44 for i in range(0,VERTICES):
     45     x = random.random() * 1920
     46     y = random.random() * 1080
     47     graph[Vertex(x,y)] = []
     48 
     49 
     50 edge_list = []
     51 keys = list(graph.keys())
     52 
     53 for i in range(0,EDGES):
     54     k1 = None
     55     k2 = None
     56     while k1 == k2:
     57         k1 = random.choice(keys)
     58         k2 = random.choice(keys)
     59 
     60     edge = Edge(k1,k2)
     61     graph[k1].append(edge)
     62     graph[k2].append(edge)
     63     edge_list.append(edge)
     64 
     65 
     66 edge_heap = []
     67 visited_vertices = set()
     68 
     69 
     70 start = random.choice(keys)
     71 start.visited = True
     72 visited_vertices.add(start)
     73 for edge in graph[start]:
     74     heapq.heappush(edge_heap, edge)
     75 
     76 
     77 first = True
     78 
     79 to_draw_vert = []
     80 to_draw_edge = []
     81 
     82 #itr = 0
     83 while True:
     84     for event in pygame.event.get():
     85         if event.type == pygame.QUIT:
     86             pygame.quit()
     87             quit()
     88 
     89     if first:
     90         display.fill(black)
     91 
     92         for edge in edge_list:
     93             edge.draw_edge(display, light_grey)
     94         first = False
     95 
     96         for vertex in graph:
     97             vertex.draw_vertex(display)
     98 
     99     for edge in to_draw_edge:
    100         edge.draw_edge(display, white)
    101     for vert in to_draw_vert:
    102         vert.draw_vertex(display)
    103 
    104     to_draw_vert = []
    105     to_draw_edge = []
    106 
    107 
    108     item = None
    109     while edge_heap:
    110         candidate = heapq.heappop(edge_heap)
    111         if (candidate.v1 in visited_vertices) != (candidate.v2 in visited_vertices):
    112             item = candidate
    113             break
    114 
    115     if item is None:
    116         break
    117 
    118     new_vertex = item.v2 if item.v1 in visited_vertices else item.v1
    119     visited_vertices.add(new_vertex)
    120     new_vertex.visited = True
    121 
    122     to_draw_vert.append(new_vertex)
    123     to_draw_edge.append(item)
    124 
    125     for edge in graph[new_vertex]:
    126         if not edge.v1 in visited_vertices or not edge.v2 in visited_vertices:
    127             heapq.heappush(edge_heap, edge)
    128     #pygame.image.save(display,"out"+str(itr)+".jpg")
    129     #itr += 1
    130 
    131     pygame.display.update()