quadtree: use internal maxHeap

quadtree/maxheap.go

@@ -0,0 +1,83 @@
+package quadtree
+
+import "github.com/paulmach/orb"
+
+// maxHeap is used for the knearest list. We need a way to maintain
+// the furthest point from the query point in the list, hence maxHeap.
+// When we find a point closer than the furthest away, we remove
+// furthest and add the new point to the heap.
+type maxHeap []*heapItem
+
+type heapItem struct {
+	point    orb.Pointer
+	distance float64
+}
+
+func (h *maxHeap) Push(item *heapItem) {
+	*h = append(*h, item)
+
+	i := len(*h) - 1
+	for i > 0 {
+		up := ((i + 1) >> 1) - 1
+		parent := (*h)[up]
+
+		if item.distance < parent.distance {
+			// parent is further so we're done fixing up the heap.
+			break
+		}
+
+		// swap nodes
+		(*h)[i] = parent
+		(*h)[up] = item
+
+		i = up
+	}
+}
+
+func (h *maxHeap) Pop() *heapItem {
+	removed := (*h)[0]
+	lastItem := (*h)[len(*h)-1]
+	(*h) = (*h)[:len(*h)-1]
+
+	mh := (*h)
+	if len(mh) == 0 {
+		return removed
+	}
+
+	// move the last item to the top and reset the heap
+	mh[0] = lastItem
+
+	i := 0
+	current := mh[i]
+	for {
+		right := (i + 1) << 1
+		left := right - 1
+
+		childIndex := i
+		child := mh[childIndex]
+
+		// swap with biggest child
+		if left < len(mh) && child.distance < mh[left].distance {
+			childIndex = left
+			child = mh[left]
+		}
+
+		if right < len(mh) && child.distance < mh[right].distance {
+			childIndex = right
+			child = mh[right]
+		}
+
+		// non bigger, so quit
+		if childIndex == i {
+			break
+		}
+
+		// swap the nodes
+		mh[i] = child
+		mh[childIndex] = current
+
+		i = childIndex
+	}
+
+	return removed
+}

quadtree/maxheap_test.go

@@ -0,0 +1,27 @@
+package quadtree
+
+import (
+	"math/rand"
+	"testing"
+)
+
+func TestMaxHeap(t *testing.T) {
+	r := rand.New(rand.NewSource(22))
+
+	for i := 1; i < 100; i++ {
+		h := make(maxHeap, 0, i)
+		for j := 0; j < i; j++ {
+			h.Push(&heapItem{distance: r.Float64()})
+		}
+
+		current := h.Pop().distance
+		for len(h) > 0 {
+			next := h.Pop().distance
+			if next > current {
+				t.Errorf("incorrect")
+			}
+
+			current = next
+		}
+	}
+}

quadtree/quadtree.go

@@ -5,7 +5,6 @@
 package quadtree
 
 import (
-	"container/heap"
 	"errors"
 	"math"
 
@@ -233,7 +232,7 @@ func (q *Quadtree) KNearestMatching(buf []orb.Pointer, p orb.Point, k int, f Fil
 		point:          p,
 		filter:         f,
 		k:              k,
-		closest:        newPointsQueue(k),
+		maxHeap:        make(maxHeap, 0, k),
 		closestBound:   &b,
 		maxDistSquared: math.MaxFloat64,
 	}
@@ -250,19 +249,14 @@ func (q *Quadtree) KNearestMatching(buf []orb.Pointer, p orb.Point, k int, f Fil
 	)
 
 	//repack result
-	if cap(buf) < len(v.closest) {
-		buf = make([]orb.Pointer, len(v.closest))
+	if cap(buf) < len(v.maxHeap) {
+		buf = make([]orb.Pointer, len(v.maxHeap))
 	} else {
-		buf = buf[:len(v.closest)]
+		buf = buf[:len(v.maxHeap)]
 	}
 
-	for i := len(v.closest) - 1; i >= 0; i-- {
-		// Actually this is a hack. We know how heap works and obtain
-		// top element without function call
-		top := v.closest[0]
-		buf[i] = top.point
-
-		heap.Pop(&v.closest)
+	for i := len(v.maxHeap) - 1; i >= 0; i-- {
+		buf[i] = v.maxHeap.Pop().point
 	}
 
 	return buf
@@ -411,53 +405,53 @@ func (v *findVisitor) Visit(n *node) {
 	}
 }
 
-type pointsQueueItem struct {
-	point    orb.Pointer
-	distance float64 // distance to point and priority inside the queue
-	index    int     // point index in queue
-}
+// type pointsQueueItem struct {
+// 	point    orb.Pointer
+// 	distance float64 // distance to point and priority inside the queue
+// 	index    int     // point index in queue
+// }
 
-type pointsQueue []pointsQueueItem
+// type pointsQueue []pointsQueueItem
 
-func newPointsQueue(capacity int) pointsQueue {
-	// We make capacity+1 because we need additional place for the greatest element
-	return make([]pointsQueueItem, 0, capacity+1)
-}
+// func newPointsQueue(capacity int) pointsQueue {
+// 	// We make capacity+1 because we need additional place for the greatest element
+// 	return make([]pointsQueueItem, 0, capacity+1)
+// }
 
-func (pq pointsQueue) Len() int { return len(pq) }
+// func (pq pointsQueue) Len() int { return len(pq) }
 
-func (pq pointsQueue) Less(i, j int) bool {
-	// We want pop longest distances so Less was inverted
-	return pq[i].distance > pq[j].distance
-}
+// func (pq pointsQueue) Less(i, j int) bool {
+// 	// We want pop longest distances so Less was inverted
+// 	return pq[i].distance > pq[j].distance
+// }
 
-func (pq pointsQueue) Swap(i, j int) {
-	pq[i], pq[j] = pq[j], pq[i]
-	pq[i].index = i
-	pq[j].index = j
-}
+// func (pq pointsQueue) Swap(i, j int) {
+// 	pq[i], pq[j] = pq[j], pq[i]
+// 	pq[i].index = i
+// 	pq[j].index = j
+// }
 
-func (pq *pointsQueue) Push(x interface{}) {
-	n := len(*pq)
-	item := x.(pointsQueueItem)
-	item.index = n
-	*pq = append(*pq, item)
-}
+// func (pq *pointsQueue) Push(x interface{}) {
+// 	n := len(*pq)
+// 	item := x.(pointsQueueItem)
+// 	item.index = n
+// 	*pq = append(*pq, item)
+// }
 
-func (pq *pointsQueue) Pop() interface{} {
-	old := *pq
-	n := len(old)
-	item := old[n-1]
-	item.index = -1
-	*pq = old[0 : n-1]
-	return item
-}
+// func (pq *pointsQueue) Pop() interface{} {
+// 	old := *pq
+// 	n := len(old)
+// 	item := old[n-1]
+// 	item.index = -1
+// 	*pq = old[0 : n-1]
+// 	return item
+// }
 
 type nearestVisitor struct {
 	point          orb.Point
 	filter         FilterFunc
 	k              int
-	closest        pointsQueue
+	maxHeap        maxHeap
 	closestBound   *orb.Bound
 	maxDistSquared float64
 }
@@ -478,13 +472,14 @@ func (v *nearestVisitor) Visit(n *node) {
 
 	point := n.Value.Point()
 	if d := planar.DistanceSquared(point, v.point); d < v.maxDistSquared {
-		heap.Push(&v.closest, pointsQueueItem{point: n.Value, distance: d})
-		if v.closest.Len() > v.k {
-			heap.Pop(&v.closest)
+		v.maxHeap.Push(&heapItem{point: n.Value, distance: d})
+		if len(v.maxHeap) > v.k {
+
+			v.maxHeap.Pop()
 
 			// Actually this is a hack. We know how heap works and obtain
 			// top element without function call
-			top := v.closest[0]
+			top := v.maxHeap[0]
 
 			v.maxDistSquared = top.distance