blog

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

prim-rlib.py (2851B)

      1 import pyray as pr
      2 import heapq
      3 import random
      4 import math
      5 
      6 VERTICES = 10000
      7 EDGES = 100000
      8 RADIUS = 5
      9 
     10 pr.init_window(5120,1440, "prim")
     11 
     12 class Vertex():
     13     def __init__(self, x, y):
     14         self.x = x
     15         self.y = y
     16         self.visited = False
     17     def draw_vertex(self):
     18         if self.visited:
     19             pr.draw_circle(int(self.x),int(self.y), RADIUS, pr.WHITE)
     20         else:
     21             pr.draw_circle(int(self.x),int(self.y), RADIUS, pr.GRAY)
     22 class Edge():
     23     def __init__(self, v1, v2):
     24         self.v1 = v1
     25         self.v2 = v2
     26         self.dist = math.sqrt(((v1.x - v2.x) ** 2) + ((v1.y - v2.y) ** 2))
     27     def draw_edge(self, c):
     28         pr.draw_line(int(self.v1.x), int(self.v1.y), int(self.v2.x), int(self.v2.y), c)
     29 
     30     def __lt__(self,otr):
     31         return self.dist < otr.dist
     32 
     33 
     34 graph = {}
     35 
     36 for i in range(0,VERTICES):
     37     x = random.random() * 5120
     38     y = random.random() * 1440
     39     graph[Vertex(x,y)] = []
     40 
     41 
     42 edge_list = []
     43 keys = list(graph.keys())
     44 
     45 for i in range(0,EDGES):
     46     k1 = None
     47     k2 = None
     48     while k1 == k2:
     49         k1 = random.choice(keys)
     50         k2 = random.choice(keys)
     51 
     52     edge = Edge(k1,k2)
     53     graph[k1].append(edge)
     54     graph[k2].append(edge)
     55     edge_list.append(edge)
     56 
     57 
     58 edge_heap = []
     59 visited_vertices = set()
     60 
     61 
     62 start = random.choice(keys)
     63 start.visited = True
     64 visited_vertices.add(start)
     65 for edge in graph[start]:
     66     heapq.heappush(edge_heap, edge)
     67 
     68 
     69 first = True
     70 
     71 to_draw_vert = []
     72 to_draw_edge = []
     73 
     74 texture = pr.load_render_texture(5120,1440)
     75 
     76 while True:
     77 
     78     pr.begin_texture_mode(texture)
     79 
     80     if first:
     81         pr.clear_background(pr.BLACK)
     82 
     83         for edge in edge_list:
     84             edge.draw_edge(pr.DARKGRAY)
     85         first = False
     86 
     87         for vertex in graph:
     88             vertex.draw_vertex()
     89 
     90     for edge in to_draw_edge:
     91         edge.draw_edge(pr.WHITE)
     92     for vert in to_draw_vert:
     93         vert.draw_vertex()
     94 
     95     to_draw_vert = []
     96     to_draw_edge = []
     97 
     98 
     99     item = None
    100     while edge_heap:
    101         candidate = heapq.heappop(edge_heap)
    102         if (candidate.v1 in visited_vertices) != (candidate.v2 in visited_vertices):
    103             item = candidate
    104             break
    105 
    106     if item is None:
    107         break
    108 
    109     new_vertex = item.v2 if item.v1 in visited_vertices else item.v1
    110     visited_vertices.add(new_vertex)
    111     new_vertex.visited = True
    112 
    113     to_draw_vert.append(new_vertex)
    114     to_draw_edge.append(item)
    115 
    116     for edge in graph[new_vertex]:
    117         if not edge.v1 in visited_vertices or not edge.v2 in visited_vertices:
    118             heapq.heappush(edge_heap, edge)
    119 
    120     pr.end_texture_mode()
    121 
    122     source_rec = pr.Rectangle(0, 0, 5120, 1440)
    123     dest_rec = pr.Rectangle(0, 0, 5120,1440)
    124 
    125     pr.begin_drawing()
    126     pr.draw_texture_pro(texture.texture, source_rec, dest_rec, pr.Vector2(0, 0), 0, pr.WHITE)
    127     pr.end_drawing()
    128 
    129 pr.unload_render_texture(texture)
    130 pr.close_window()