spec/kd_expected_out.txt spec/kd_test_in.txt @@ -9,11 +9,14 @@ First, put the points into the tree: - kd = Containers::KDTree.new([ [4, 3], [3, 4], [-1, 2], [6, 4], [3, -5], [-2, -5] ]) + kdtree = Containers::KDTree.new( {0 => [4, 3], 1 => [3, 4], 2 => [-1, 2], 3 => [6, 4], + 4 => [3, -5], 5 => [-2, -5] }) Then, query on the tree: - puts kd.find_nearest([0, 0], 2) => [[0, 6], [0, 3]] + puts kd.find_nearest([0, 0], 2) => [[5, 2], [9, 1]] + + The result is an array of [distance, id] pairs. There seems to be a bug in this version. Note that the point queried on does not have to exist in the tree. However, if it does exist, it will be returned. @@ -23,21 +26,30 @@ class Containers::KDTree Node = Struct.new(:id, :coords, :left, :right) + # Points is a hash of id => [coord, coord] pairs. def initialize(points) - @root = build_tree(points) + raise "must pass in a hash" unless points.kind_of?(Hash) + @dimensions = points[ points.keys.first ].size + @root = build_tree(points.to_a) + @nearest = [] + end + + # Find k closest points to given coordinates + def find_nearest(target, k_nearest) @nearest = [] + nearest(@root, target, k_nearest, 0) end - # Build a kd-tree + # points is an array def build_tree(points, depth=0) return if points.empty? - - axis = depth % 2 - - points.sort! { |a, b| a[1][axis] <=> b[1][axis] } + + axis = depth % @dimensions + + points.sort! { |a, b| a.last[axis] <=> b.last[axis] } median = points.size / 2 - - node = Node.new(points[median][0], points[median][1], nil, nil) + + node = Node.new(points[median].first, points[median].last, nil, nil) node.left = build_tree(points[0...median], depth+1) node.right = build_tree(points[median+1..-1], depth+1) node @@ -65,14 +77,9 @@ class Containers::KDTree end private :check_nearest - # Find k closest points to given coordinates - def find_nearest(target, k_nearest) - @nearest = [] - nearest(@root, target, k_nearest, 0) - end - + # Recursively find nearest coordinates, going down the appropriate branch as needed def nearest(node, target, k_nearest, depth) - axis = depth % 2 + axis = depth % @dimensions if node.left.nil? && node.right.nil? # Leaf node @nearest = check_nearest(@nearest, node, target, k_nearest) lib/containers/kd_tree.rb @@ -1,4 +1,34 @@ $: << File.join(File.expand_path(File.dirname(__FILE__)), '..', 'lib') require 'algorithms' -# TODO: KD Tree Spec +describe Containers::KDTree do + it "should work for a documented example" do + kdtree = Containers::KDTree.new( {0 => [4, 3], 1 => [3, 0], 2 => [-1, 2], 3 => [6, 4], + 4 => [3, -5], 5 => [-2, -5] }) + closest_2 = kdtree.find_nearest([0, 0], 2) + closest_2.should eql([[5, 2], [9, 1]]) + end + + it "should work for real-life example from facebook puzzle" do + points = {} + input = File.open(File.join(File.dirname(__FILE__), 'kd_test_in.txt'), 'r') + + # Populate points hash + input.each_line do |line| + break if line.empty? + n, x, y = line.split(/\s+/) + points[n.to_i] = [x.to_f, y.to_f] + end + + out = "" + kdtree = Containers::KDTree.new(points) + points.sort{ |(k1, v1), (k2, v2)| k1 <=> k2 }.each { |id, point| + nearest_4 = kdtree.find_nearest(point, 4) + out << "#{id} #{nearest_4[1..-1].collect{ |n| n[1] }.join(',')}\n" + } + + expected = File.read(File.join(File.dirname(__FILE__), 'kd_expected_out.txt')) + expected.should eql(out) + end + +end spec/kd_tree_spec.rb e26bd6a481d21f39721c2538f223459f90ae8d46 Kanwei Li kanwei@gmail.com http://github.com/kanwei/algorithms/commit/f2db2fee5a7f13000ac1f8ff24fcaed2f7283aea f2db2fee5a7f13000ac1f8ff24fcaed2f7283aea 2009-08-25T23:11:06-07:00 2009-08-25T23:11:06-07:00 Tests for kd-tree, it still seems semi-broken ec06548a59ed6aabde92701d5547e7b5653494a7 Kanwei Li kanwei@gmail.com