diff --git a/quadtree/benchmarks_test.go b/quadtree/benchmarks_test.go index 1b00aa8..908c486 100644 --- a/quadtree/benchmarks_test.go +++ b/quadtree/benchmarks_test.go @@ -132,6 +132,23 @@ func BenchmarkRandomInBound1000Buf(b *testing.B) { } } +func BenchmarkRandomKNearest10(b *testing.B) { + r := rand.New(rand.NewSource(43)) + + qt := New(orb.Bound{Min: orb.Point{0, 0}, Max: orb.Point{1, 1}}) + for i := 0; i < 1000; i++ { + qt.Add(orb.Point{r.Float64(), r.Float64()}) + } + + buf := make([]orb.Pointer, 0, 10) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + qt.KNearest(buf[:0], orb.Point{r.Float64(), r.Float64()}, 10) + } +} + func BenchmarkRandomKNearest100(b *testing.B) { r := rand.New(rand.NewSource(43)) diff --git a/quadtree/maxheap.go b/quadtree/maxheap.go index 9578150..5473a3e 100644 --- a/quadtree/maxheap.go +++ b/quadtree/maxheap.go @@ -13,27 +13,46 @@ type heapItem struct { distance float64 } -func (h *maxHeap) Push(item *heapItem) { - *h = append(*h, item) +func (h *maxHeap) Push(point orb.Pointer, distance float64) { + // Common usage is Push followed by a Pop if we have > k points. + // We're reusing the k+1 heapItem object to reduce memory allocations. + // First we manaully lengthen the slice, + // then we see if the last item has been allocated already. + + prevLen := len(*h) + *h = (*h)[:prevLen+1] + if (*h)[prevLen] == nil { + (*h)[prevLen] = &heapItem{point: point, distance: distance} + } else { + (*h)[prevLen].point = point + (*h)[prevLen].distance = distance + } i := len(*h) - 1 for i > 0 { up := ((i + 1) >> 1) - 1 parent := (*h)[up] - if item.distance < parent.distance { + if distance < parent.distance { // parent is further so we're done fixing up the heap. break } // swap nodes - (*h)[i] = parent - (*h)[up] = item + // (*h)[i] = parent + (*h)[i].point = parent.point + (*h)[i].distance = parent.distance + + // (*h)[up] = item + (*h)[up].point = point + (*h)[up].distance = distance i = up } } +// Pop returns the "greatest" item in the list. +// The returned item should not be saved across push/pop operations. func (h *maxHeap) Pop() *heapItem { removed := (*h)[0] lastItem := (*h)[len(*h)-1] diff --git a/quadtree/maxheap_test.go b/quadtree/maxheap_test.go index 9a1df34..885438a 100644 --- a/quadtree/maxheap_test.go +++ b/quadtree/maxheap_test.go @@ -11,7 +11,7 @@ func TestMaxHeap(t *testing.T) { for i := 1; i < 100; i++ { h := make(maxHeap, 0, i) for j := 0; j < i; j++ { - h.Push(&heapItem{distance: r.Float64()}) + h.Push(nil, r.Float64()) } current := h.Pop().distance diff --git a/quadtree/quadtree.go b/quadtree/quadtree.go index 843af16..5b6b28b 100644 --- a/quadtree/quadtree.go +++ b/quadtree/quadtree.go @@ -212,6 +212,7 @@ func (q *Quadtree) Matching(p orb.Point, f FilterFunc) orb.Pointer { // KNearest returns k closest Value/Pointer in the quadtree. // This function is thread safe. Multiple goroutines can read from a pre-created tree. // An optional buffer parameter is provided to allow for the reuse of result slice memory. +// The points are returned in a sorted order, nearest first. // This function allows defining a maximum distance in order to reduce search iterations. func (q *Quadtree) KNearest(buf []orb.Pointer, p orb.Point, k int, maxDistance ...float64) []orb.Pointer { return q.KNearestMatching(buf, p, k, nil, maxDistance...) @@ -221,6 +222,7 @@ func (q *Quadtree) KNearest(buf []orb.Pointer, p orb.Point, k int, maxDistance . // the given filter function returns true. This function is thread safe. // Multiple goroutines can read from a pre-created tree. An optional buffer // parameter is provided to allow for the reuse of result slice memory. +// The points are returned in a sorted order, nearest first. // This function allows defining a maximum distance in order to reduce search iterations. func (q *Quadtree) KNearestMatching(buf []orb.Pointer, p orb.Point, k int, f FilterFunc, maxDistance ...float64) []orb.Pointer { if q.root == nil { @@ -232,7 +234,7 @@ func (q *Quadtree) KNearestMatching(buf []orb.Pointer, p orb.Point, k int, f Fil point: p, filter: f, k: k, - maxHeap: make(maxHeap, 0, k), + maxHeap: make(maxHeap, 0, k+1), closestBound: &b, maxDistSquared: math.MaxFloat64, } @@ -472,7 +474,7 @@ func (v *nearestVisitor) Visit(n *node) { point := n.Value.Point() if d := planar.DistanceSquared(point, v.point); d < v.maxDistSquared { - v.maxHeap.Push(&heapItem{point: n.Value, distance: d}) + v.maxHeap.Push(n.Value, d) if len(v.maxHeap) > v.k { v.maxHeap.Pop()