Skip to content

Instantly share code, notes, and snippets.

@648trindade
Created March 4, 2020 19:45
Show Gist options
  • Save 648trindade/62edf808dec6b31a44422adf89eec780 to your computer and use it in GitHub Desktop.
Save 648trindade/62edf808dec6b31a44422adf89eec780 to your computer and use it in GitHub Desktop.
KD Tree in Pygame
import pygame, random
class Point:
def __init__(self, x, y):
self.x = x
self.y = y
class Kd_tree:
class Axis:
X = 0
Y = 1
DIMS = 2
class Leaf:
def __init__(self, point):
self.point = point
def draw(self, surface, box):
pygame.draw.circle(surface, [0,0,0], [self.point.x, self.point.y], 4)
class Node:
def __init__(self, axis, point_list, parent = None):
self.axis = axis
self.value = None
self.parent = parent
mid = None
if self.axis is Kd_tree.Axis.X:
point_list = sorted(point_list, key=lambda p: p.x)
else:
point_list = sorted(point_list, key=lambda p : p.y)
mid = len(point_list) // 2
if len(point_list) > 2:
if self.axis is Kd_tree.Axis.X:
self.value = point_list[mid-1].x
else:
self.value = point_list[mid-1].y
self.left = Kd_tree.Node((axis+1) % Kd_tree.Axis.DIMS, point_list[:mid], self)
self.right = Kd_tree.Node((axis+1) % Kd_tree.Axis.DIMS, point_list[mid:], self)
else:
self.left = Kd_tree.Leaf(point_list[0])
self.right = None
if self.axis is Kd_tree.Axis.X:
self.value = point_list[0].x
else:
self.value = point_list[0].y
if len(point_list) == 2:
self.right = Kd_tree.Leaf(point_list[1])
def draw(self, surface, box):
if self.axis is Kd_tree.Axis.X:
pygame.draw.line(surface, [255,0,0], [self.value, box[0][1]], [self.value, box[1][1]])
self.left.draw(surface, [
box[0],
[self.value, box[1][1]]
])
if self.right:
self.right.draw(surface, [
[self.value, box[0][1]],
box[1]
])
else:
pygame.draw.line(surface, [0,255,0], [box[0][0], self.value], [box[1][0], self.value])
self.left.draw(surface, [
box[0],
[box[0][0], self.value]
])
if self.right:
self.right.draw(surface, [
[box[0][0], self.value],
box[1]
])
def __init__(self, point_list):
self.root = self.Node(Kd_tree.Axis.X, point_list)
def draw(self, surface, box):
self.root.draw(surface, box)
points = [Point(random.randint(0,800), random.randint(0,600)) for i in range(10)]
tree = Kd_tree(points)
pygame.init()
surface = pygame.display.set_mode([800, 600])
while True:
surface.fill([255,255,255])
tree.draw(surface, [[0,0], [800,600]])
pygame.display.flip()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment