Skip to content

Instantly share code, notes, and snippets.

@tscolari
Created December 14, 2012 20:15
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save tscolari/4288299 to your computer and use it in GitHub Desktop.
Save tscolari/4288299 to your computer and use it in GitHub Desktop.
Simple KDTree implementation to find the nearest point to another
# Simple KDTree class with sorting for finding nearest neighbor
#
class KDTree
class Node < Struct.new(:value, :left, :right)
# Public: Returns the distance of the object to the
# point
#
def distance_to(point)
throw "Points have different dimensions" if point.size != value.size
square_sum = 0
point.each_index do |i|
square_sum += (point[i] - value[i])**2
end
Math.sqrt(square_sum)
end
# Public: Goes recursively through nodes childs to find
# wich one is closest to the point.
# returns an hash with:
# :node => the node element that is closest
# :distance => distance of this node element to the point
#
def find_closest_child_to(point, dimension = 0)
my_distance = distance_to(point)
next_dimension = KDTree.toggle_dimension(dimension, value.length)
if value[dimension] > point[dimension]
if self.left
child_distance = self.left.find_closest_child_to(point, next_dimension)
return child_distance if child_distance[:distance] < my_distance
end
else
if self.right
child_distance = self.right.find_closest_child_to(point, next_dimension)
return child_distance if child_distance[:distance] < my_distance
end
end
{ point: self, distance: my_distance }
end
end
def initialize(space)
@root = build_node_for(space, 0)
end
def build_node_for(space, dimension = 0)
return nil if space.empty?
sorted_space = sort_space_by_dimension(space, dimension)
space_values = space_partition(sorted_space)
next_dimension = KDTree.toggle_dimension(dimension, space_values[:value].length)
Node.new(
space_values[:value],
build_node_for(space_values[:left] , next_dimension) ,
build_node_for(space_values[:right] , next_dimension)
)
end
def nearest_node_to(point)
@root.find_closest_child_to(point)[:point]
end
def self.toggle_dimension(dimension, max_dimensions = 3)
(dimension + 1) % max_dimensions
end
private
# Private: Breaks the space into a hash that has:
# :value => the middle point of the space
# :left => the points that are left from the 'value' in the array
# :right => the points that are right from the 'value' in the array
#
def space_partition(space)
left_values = space.first((space.length / 2.0).ceil)
right_values = space - left_values
mid_value = left_values.delete_at(left_values.length - 1)
{ value: mid_value, left: left_values, right: right_values }
end
# Private: Sorts the space array based in one dimension
#
def sort_space_by_dimension(space, dimension = 0)
space.sort! do |a,b|
a[dimension] <=> b[dimension]
end
end
end
require './kd_tree'
describe KDTree do
describe "Node" do
describe "#distance_to" do
[
{point1: [0,0], point2: [0,1], distance: 1},
{point1: [0,4], point2: [3,0], distance: 5},
{point1: [-2,1], point2: [1,5], distance: 5}
].each do |points|
it "should return #{points[:distance]} for #{points[:point1]} and #{points[:point2]}" do
node = KDTree::Node.new(points[:point1], nil, nil)
node.distance_to(points[:point2]).should == points[:distance]
end
end
end
describe "#find_closest_child_to" do
let(:right_node) { KDTree::Node.new([100,100], nil, nil) }
let(:left_node) { KDTree::Node.new([-100,-100], nil, nil) }
let(:node) { KDTree::Node.new([0,0], left_node, right_node) }
context "node is a leaf" do
it "should return itself" do
node = KDTree::Node.new([0,0], nil, nil)
node.find_closest_child_to([1,1],0).should == { point: node, distance: node.distance_to([1,1])}
end
end
context "should go left" do
it "should return the left node" do
node.find_closest_child_to([-90,-90],0).should == { point: left_node, distance: left_node.distance_to([-90, -90])}
end
it "should return the root node" do
right_node.should_receive(:find_closest_child_to).never
node.find_closest_child_to([-5,-5],0).should == { point: node, distance: node.distance_to([-5, -5])}
end
end
context "should go right" do
it "should return the right node" do
node.find_closest_child_to([90,90],0).should == { point: right_node, distance: right_node.distance_to([90,90])}
end
it "should return the root node" do
left_node.should_receive(:find_closest_child_to).never
node.find_closest_child_to([5,5],0).should == { point: node, distance: node.distance_to([5, 5])}
end
end
end
end # "Node"
describe "#toggle_dimension" do
it "should cycle the dimension" do
KDTree.toggle_dimension(0, 3).should eq 1
KDTree.toggle_dimension(1, 3).should eq 2
KDTree.toggle_dimension(2, 3).should eq 0
KDTree.toggle_dimension(3, 3).should eq 1
KDTree.toggle_dimension(0, 2).should eq 1
KDTree.toggle_dimension(1, 2).should eq 0
KDTree.toggle_dimension(2, 2).should eq 1
end
end
describe "#space_partition" do
[
{
space: [[-2,1],[1,1],[2,2],[3,3],[10,50]],
result: {value: [2,2], left: [[-2,1],[1,1]], right: [[3,3],[10,50]]}
},
{
space: [[3,3],[10,50]],
result: {value: [3,3], left: [], right: [[10,50]]}
},
{
space: [[-2,1],[1,1]],
result: {value: [-2,1], left: [], right: [[1,1]]}
},
{
space: [[-2,1],[0,0],[1,1]],
result: {value: [0,0], left: [[-2,1]], right: [[1,1]]}
}
].each do |example|
it "should split '#{example[:space]}' in a hash with the 3 correct keys" do
KDTree.any_instance.stub(:build_node_for).and_return(nil)
kd_tree = KDTree.new([])
kd_tree.send(:space_partition, example[:space]).should eq example[:result]
end
end
end
describe "#sort_space_by_dimension" do
[
{
dimension: 0,
original: [[3,3],[1,1],[2,2],[-2,5],[10,50]],
sorted: [[-2,5],[1,1],[2,2],[3,3],[10,50]]
},
{
dimension: 1,
original: [[3,3],[1,1],[2,2],[-2,5],[10,50]],
sorted: [[1,1],[2,2],[3,3],[-2,5],[10,50]]
}
].each do |example|
it "should sort '#{example[:original]}' by dimension #{example[:dimension]} correctly" do
KDTree.any_instance.stub(:build_node_for).and_return(nil)
kd_tree = KDTree.new([])
kd_tree.send(:sort_space_by_dimension, example[:original], example[:dimension]).should eq example[:sorted]
end
end
end
describe "#nearest_node_to" do
[
{
space: [[3,3],[1,1],[2,2],[-2,5], [10,50], [-10, 20]],
tests: [
{ point: [0,0], correct: [1,1] },
{ point: [2,2], correct: [2,2] },
{ point: [100,100], correct: [10,50] },
{ point: [30,-10], correct: [3,3] }
]
},
].each do |example|
context "Space: #{example[:space]}" do
it "should return the correct nodes" do
tree = KDTree.new(example[:space])
example[:tests].each do |test|
tree.nearest_node_to(test[:point]).value.should eq test[:correct]
end
end
end
end
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment