Simple KDTree implementation to find the nearest point to another

kd_tree.rb

# 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

kd_tree_spec.rb

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