kd_tree.rb

Node = Struct.new(:value, :left, :right)

class KDTree
  attr_reader :root
  
  # points - an array of arrays containing x,y coordinates (e.g. [[1,1], [5,1]])
  def initialize(points)
    @root = init_tree(points.clone)
  end
  
  def init_tree(points)
    axis = points.length % 2
    points = points.sort_by { |point| point[axis] }
    median = points.length / 2
    
    if median == 0
      return Node.new(points[median], nil, nil)
    end
    
    Node.new(points[median],
             init_tree(points[0..median - 1]),
             init_tree(points[median + 1..-1]))
  end
end