Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Build a kd-tree for 2d points, with two consecutive x-levels and one y-level.
#!/usr/bin/python
import sys
import numpy
import Queue
import numpy as np
import random
import matplotlib.pyplot as plt
# Generate set of points
S = [
[11, 0],
[15, 4],
[10, 13],
[5, 12],
[4, 19],
[0, 17],
[8, 30],
[10, 23],
[9, 17],
[6, 27],
[8, 23],
[11, 8]
]
X, Y = zip(*S)
# Render all points
plt.plot(X, Y, 'ro')
maxX = max(X)
maxY = max(Y)
print maxX
class Node:
l = 0
left = None
right = None
level = 0
def __init__(self, l, left, right, level):
self.l = l
self.left = left
self.right = right
self.level = level
class Leaf:
v = 0
level = 0
def __init__(self, v, level):
self.v = v
self.level = level
def median(lst):
return numpy.median(numpy.array(lst))
def buildKdTree(P, depth):
if len(P) == 1:
return Leaf(P[0], depth)
else:
# Split the set in X and Y coordinates
X, Y = zip(*P)
# Y level
if depth % 3 == 0:
# Calculate the median
l = median(Y)
# Split the set of points
P1 = [p for p in P if p[1] <= l]
P2 = [p for p in P if p[1] > l]
# Render a line
plt.plot([0, maxX], [l, l])
# X level
else:
# Calculate the median
l = median(X)
# Split the set of points
P1 = [p for p in P if p[0] <= l]
P2 = [p for p in P if p[0] > l]
# Render a line
plt.plot([l, l], [0, maxY])
# Construct a left and right child
left = buildKdTree(P1, depth + 1)
right = buildKdTree(P2, depth + 1)
node = Node(l, left, right, depth)
return node
def displayTree(tree):
# Create a queue
queue = Queue.Queue()
# Add the root
queue.put(tree)
# Keep track of the last level
lastLevel = 1
# Display
while not queue.empty():
# Retrieve the current
current = queue.get()
# Check if we should add a new line
if current.level > lastLevel:
sys.stdout.write("\n")
lastLevel += 1
# Display data of current
if isinstance(current, Node):
sys.stdout.write('N(%f)' % current.l)
# Recurse
if current.left != None:
queue.put(current.left)
if current.right != None:
queue.put(current.right)
elif isinstance(current, Leaf):
sys.stdout.write('L(%i,%i)' % (current.v[0], current.v[1]))
# Build the 2d tree
tree = buildKdTree(S, 1)
# Display the tree
displayTree(tree)
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment