diff --git a/README.md b/README.md index 320fc0f..0f2f3d4 100644 --- a/README.md +++ b/README.md @@ -73,15 +73,15 @@ Using a tree with 1 million points `[x, y] of Float64` on my i7-8550U CPU @ 1.80 ```console Benchmarking KD-Tree with 1 million points -build(init): 3.41 seconds +build(init): 4.34 seconds user system total real -nearest point 1 0.000019 0.000000 0.000019 ( 0.000019) -nearest point 5 0.000021 0.000000 0.000021 ( 0.000021) -nearest point 10 0.000025 0.000001 0.000026 ( 0.000025) -nearest point 50 0.000269 0.000002 0.000271 ( 0.000272) -nearest point 100 0.000809 0.000000 0.000809 ( 0.000812) -nearest point 255 0.005078 0.000000 0.005078 ( 0.005087) -nearest point 999 0.439598 0.000001 0.439599 ( 0.440699) +nearest point 1 0.000017 0.000001 0.000018 ( 0.000017) +nearest point 5 0.000022 0.000000 0.000022 ( 0.000022) +nearest point 10 0.000021 0.000001 0.000022 ( 0.000022) +nearest point 50 0.000058 0.000001 0.000059 ( 0.000059) +nearest point 100 0.000087 0.000002 0.000089 ( 0.000089) +nearest point 255 0.000248 0.000005 0.000253 ( 0.000254) +nearest point 999 0.001033 0.000020 0.001053 ( 0.001055) ``` ## Contributing diff --git a/shard.yml b/shard.yml index 1b61377..446cff8 100644 --- a/shard.yml +++ b/shard.yml @@ -7,6 +7,11 @@ description: | authors: - Anton Maminov +dependencies: + priority-queue: + github: spider-gazelle/priority-queue + branch: master + development_dependencies: ameba: github: crystal-ameba/ameba diff --git a/src/ext/priority-queue.cr b/src/ext/priority-queue.cr new file mode 100644 index 0000000..712d435 --- /dev/null +++ b/src/ext/priority-queue.cr @@ -0,0 +1,16 @@ +require "priority-queue" + +module Priority + class Item(V) + def initialize(@priority : Float64, @value : V, name = nil) + @name = name.to_s if name + end + end + + class Queue(V) + def push(priority : Float64, value : V, name = nil) + item = Item(V).new(priority, value, name) + push(item) + end + end +end diff --git a/src/kd_tree.cr b/src/kd_tree.cr index 1f920ba..334daa6 100644 --- a/src/kd_tree.cr +++ b/src/kd_tree.cr @@ -1,3 +1,4 @@ +require "./ext/priority-queue" require "./kd_tree/*" module Kd @@ -25,18 +26,18 @@ module Kd def nearest(target : Array(T), n : Int32 = 1) : Array(Array(T)) return [] of Array(T) if n < 1 - best_nodes = Array(Node(T)).new + best_nodes = Priority::Queue(Node(T)).new find_n_nearest(@root, target, 0, best_nodes, n) - best_nodes.map(&.pivot) + best_nodes.map(&.value.pivot) end private def find_n_nearest( node : Node(T)?, target : Array(T), depth : Int32, - best_nodes : Array(Node(T)), + best_nodes : Priority::Queue(Node(T)), n : Int32 ) return unless node @@ -48,29 +49,25 @@ module Kd find_n_nearest(next_node, target, depth + 1, best_nodes, n) - if best_nodes.size < n || distance(target, node.pivot) < distance(target, best_nodes.last.pivot) - best_nodes << node - best_nodes.sort_by! { |nd| distance(target, nd.pivot) } - best_nodes.pop if best_nodes.size > n - end + best_nodes.push(distance(target, node.pivot), node) + + best_nodes.pop if best_nodes.size > n - if other_node && (best_nodes.size < n || (target[axis] - node.pivot[axis]).abs**2 < distance(target, best_nodes.last.pivot)) + if other_node && (best_nodes.size < n || (target[axis] - node.pivot[axis]).abs ** 2 < distance(target, best_nodes.last.value.pivot)) find_n_nearest(other_node, target, depth + 1, best_nodes, n) end end private def distance(m : Array(T), n : Array(T)) # squared euclidean distance (to avoid expensive sqrt operation) - m.each_with_index.sum do |coord, index| - (coord - n[index]) ** 2 - end + m.each_with_index.sum { |coord, index| (coord - n[index]) ** 2 } end private def build_tree(points : Array(Array(T)), depth : Int32) : Node(T)? return if points.empty? axis = depth % @k - points.sort_by! { |point| point[axis] } + points.sort_by!(&.[axis]) median = points.size // 2 # Create node and construct subtrees