blog

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

prim.py (2694B)

      1 import pygame
      2 import heapq
      3 import random
      4 import math
      5 
      6 pygame.init()
      7 
      8 display = pygame.display.set_mode((5120,1440))
      9 
     10 VERTICES = 1000
     11 EDGES = 10000
     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 
     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 
     32 class Edge():
     33     def __init__(self, v1, v2):
     34         self.v1 = v1
     35         self.v2 = v2
     36         self.dist = math.sqrt(((v1.x - v2.x) ** 2) + ((v1.y - v2.y) ** 2))
     37 
     38     def draw_edge(self,display, c):
     39         pygame.draw.line(display, c, (self.v1.x,self.v1.y), (self.v2.x,self.v2.y), 1)
     40 
     41     def __lt__(self,otr):
     42         return self.dist < otr.dist
     43 
     44 
     45 graph = {}
     46 
     47 for i in range(0,VERTICES):
     48     x = random.random() * 5120
     49     y = random.random() * 1440
     50     graph[Vertex(x,y)] = []
     51 
     52 
     53 edge_list = []
     54 for i in range(0,EDGES):
     55     k1 = None
     56     k2 = None
     57     while k1 == k2:
     58         k1 = random.choice(list(graph.keys()))
     59         k2 = random.choice(list(graph.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 mst = []
     71 
     72 start = random.choice(list(graph.keys()))
     73 start.visited = True
     74 visited_vertices.add(start)
     75 for edge in graph[start]:
     76     heapq.heappush(edge_heap, edge)
     77 
     78 while True:
     79     for event in pygame.event.get():
     80         if event.type == pygame.QUIT:
     81             pygame.quit()
     82             quit()
     83 
     84     display.fill(black)
     85     for edge in edge_list:
     86         edge.draw_edge(display, light_grey)
     87     for edge in mst:
     88         edge.draw_edge(display, white)
     89     for vertex in graph:
     90         vertex.draw_vertex(display)
     91 
     92     item = None
     93     while edge_heap:
     94         candidate = heapq.heappop(edge_heap)
     95         if (candidate.v1 in visited_vertices) != (candidate.v2 in visited_vertices):
     96             item = candidate
     97             break
     98 
     99     if item is None:
    100         pygame.display.update()
    101         break
    102 
    103     new_vertex = item.v2 if item.v1 in visited_vertices else item.v1
    104     visited_vertices.add(new_vertex)
    105     new_vertex.visited = True
    106     mst.append(item)
    107     for edge in graph[new_vertex]:
    108         heapq.heappush(edge_heap, edge)
    109 
    110     #dir_name = "/dev/shm/bg/"
    111     #if not os.path.exists(dir_name):
    112     #    os.mkdir(dir_name)
    113     #pygame.image.save(display, dir_name + "out.png")
    114     #os.system("/usr/bin/feh --no-fehbg --bg-tile '/dev/shm/bg/out.png' ")
    115     #time.sleep(.1)
    116     pygame.display.update()