improved quadtree, kdtree tests and benchmarks

Author: aldernero

kdtree.go

@@ -1,6 +1,9 @@
 package sketchy
 
-import "github.com/tdewolff/canvas"
+import (
+	"container/heap"
+	"github.com/tdewolff/canvas"
+)
 
 type KDTree struct {
 	point  *Point
@@ -67,6 +70,27 @@ func (k *KDTree) Query(r Rect) []Point {
 	return results
 }
 
+func (k *KDTree) NearestNeighbors(point Point, s int) []Point {
+	var result []Point
+	if s <= 0 {
+		return result
+	}
+	ph := &PointHeap{
+		size:   s,
+		points: []MetricPoint{},
+	}
+	heap.Init(ph)
+	k.pushOnHeap(point, ph)
+	result = ph.Report()
+	return result
+}
+
+func (k *KDTree) Clear() {
+	k.point = nil
+	k.left = nil
+	k.right = nil
+}
+
 func (k *KDTree) Size() int {
 	count := 1
 	if k.left != nil {
@@ -86,12 +110,6 @@ func (k *KDTree) DrawWithPoints(s float64, ctx *canvas.Context) {
 	k.draw(ctx, 0, s)
 }
 
-func (k *KDTree) Clear() {
-	k.point = nil
-	k.left = nil
-	k.right = nil
-}
-
 func (k *KDTree) insert(p Point, d int) {
 	if k.point == nil {
 		k.point = &p
@@ -188,3 +206,18 @@ func (k *KDTree) draw(ctx *canvas.Context, depth int, pointSize float64) {
 		k.right.draw(ctx, depth+1, pointSize)
 	}
 }
+
+func (k *KDTree) pushOnHeap(target Point, h *PointHeap) {
+	if k.point != nil {
+		heap.Push(h, MetricPoint{
+			Metric: SquaredDistance(target, *k.point),
+			Point:  *k.point,
+		})
+	}
+	if k.left != nil {
+		k.left.pushOnHeap(target, h)
+	}
+	if k.right != nil {
+		k.right.pushOnHeap(target, h)
+	}
+}