Merge pull request #6 from geocrystal/priority-queue

use priority queue

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

shard.yml

@@ -7,6 +7,11 @@ description: |
 authors:
   - Anton Maminov <[email protected]>
 
+dependencies:
+  priority-queue:
+    github: spider-gazelle/priority-queue
+    branch: master
+
 development_dependencies:
   ameba:
     github: crystal-ameba/ameba

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

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