blog

Personal blog
git clone git://git.laack.co/blog.git
Log | Files | Refs

prim.py (2844B)

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