algorithms

Algorithm implementations
git clone git://git.laack.co/algorithms.git
Log | Files | Refs | README

prim.py (2805B)

      1 import pygame
      2 import time
      3 import heapq
      4 import random
      5 import math
      6 
      7 pygame.init()
      8 
      9 display = pygame.display.set_mode((5120,1440))
     10 
     11 VERTICES = 100
     12 EDGES = 1000
     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 
     27     def draw_vertex(self,display):
     28         if self.visited:
     29             pygame.draw.circle(display, white, (self.x,self.y), 5)
     30         else:
     31             pygame.draw.circle(display, grey, (self.x,self.y), 5)
     32 
     33 class Edge():
     34     def __init__(self, v1, v2):
     35         self.v1 = v1
     36         self.v2 = v2
     37         self.dist = math.sqrt(((v1.x - v2.x) ** 2) + ((v1.y - v2.y) ** 2))
     38 
     39     def draw_edge(self,display, c):
     40         pygame.draw.line(display, c, (self.v1.x,self.v1.y), (self.v2.x,self.v2.y), 1)
     41 
     42     def __lt__(self,otr):
     43         return self.dist < otr.dist
     44 
     45 
     46 while True:
     47     graph = {}
     48 
     49     for i in range(0,VERTICES):
     50         x = random.random() * 5120
     51         y = random.random() * 1440
     52         graph[Vertex(x,y)] = []
     53 
     54 
     55     edge_list = []
     56     for i in range(0,EDGES):
     57         k1 = None
     58         k2 = None
     59         while k1 == k2:
     60             k1 = random.choice(list(graph.keys()))
     61             k2 = random.choice(list(graph.keys()))
     62 
     63         edge = Edge(k1,k2)
     64         graph[k1].append(edge)
     65         graph[k2].append(edge)
     66         edge_list.append(edge)
     67 
     68 
     69     edge_heap = []
     70     visited_vertices = set()
     71 
     72     mst = []
     73 
     74     start = random.choice(list(graph.keys()))
     75     start.visited = True
     76     visited_vertices.add(start)
     77     for edge in graph[start]:
     78         heapq.heappush(edge_heap, edge)
     79 
     80     while True:
     81         for event in pygame.event.get():
     82             if event.type == pygame.QUIT:
     83                 pygame.quit()
     84                 quit()
     85 
     86         display.fill(black)
     87         for vertex in graph:
     88             vertex.draw_vertex(display)
     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 
     94         item = None
     95         while edge_heap:
     96             candidate = heapq.heappop(edge_heap)
     97             if (candidate.v1 in visited_vertices) != (candidate.v2 in visited_vertices):
     98                 item = candidate
     99                 break
    100 
    101         if item is None:
    102             pygame.display.update()
    103             if len(visited_vertices) == VERTICES:
    104                 pygame.image.save(display, "out.jpg")
    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         time.sleep(.1)
    115         pygame.display.update()