blog

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

prim.py (2820B)

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