perf!: ⚡ improve the time and space efficiency of the KNearestNeighbor API

The fundamental changes made to improve the time and space efficiency of the `KNearestNeighbor` API are:
1. Replace the interface-based implementation of the priority queue with a generics-based implementation.
2. Preallocate an underlying slice with a capacity equal to the maximum allowed size of the priority queue.
3. The `Item` struct now stores a pointer to the original data instead of storing a copy of the original data.
4. Return a slice of pointers to the k-nearest neighbours (i.e., `[]*T`) instead of a slice of copies of the k-nearest neighbours' values (i.e., `[]T`).

These changes have led to a remarkable ~34.26% enhancement in performance (i.e., `ns/op`). The memory efficiency has also substantially improved, with ~5 times less `B/op` and ~113 times less `allocs/op`.

BREAKING CHANGE: To improve the performance of the `KNearestNeighbor` API, I changed its return type from `[]T` to `[]*T` to reduce the number of costly data copies that occur. Utilizing pointers to access the actual data is a significantly more efficient approach in terms of CPU and memory usage. With just a single level of indirection, we can save on a considerable amount of data copying, ensuring a more streamlined implementation. With this change to the `KNearestNeighbor` API's return type in this commit, any code that uses the `KNearestNeighbor` API will no longer work with future releases of this library created after this commit. Affected codebases must migrate to the updated version of the `KNearestNeighbor` API to be compatible with future releases of this library.

Author: rishitc

internal/utils/generic_pq.go

@@ -0,0 +1,80 @@
+// This file is a rewrite of Go's "container/heap" implementation with the usages of interfaces replaced with generics.
+// Link to Go's "container/heap" implementation: https://cs.opensource.google/go/go/+/refs/tags/go1.23.0:src/container/heap/heap.go
+
+package internal
+
+import "sort"
+
+type Interface[T any] interface {
+	sort.Interface
+	Push(x T) // add x as element Len()
+	Pop() T   // remove and return element Len() - 1.
+}
+
+func Init[V any, T Interface[V]](h T) {
+	// heapify the input
+	n := h.Len()
+	for i := n/2 - 1; i >= 0; i-- {
+		down(h, i, n)
+	}
+}
+
+func Push[V any, T Interface[V]](h T, x V) {
+	h.Push(x)
+	up(h, h.Len()-1)
+}
+
+func Pop[V any, T Interface[V]](h T) V {
+	n := h.Len() - 1
+	h.Swap(0, n)
+	down(h, 0, n)
+	return h.Pop()
+}
+
+func Remove[V any, T Interface[V]](h T, i int) any {
+	n := h.Len() - 1
+	if n != i {
+		h.Swap(i, n)
+		if !down(h, i, n) {
+			up(h, i)
+		}
+	}
+	return h.Pop()
+}
+
+func Fix[V any, T Interface[V]](h T, i int) {
+	if !down(h, i, h.Len()) {
+		up(h, i)
+	}
+}
+
+func up[V any, T Interface[V]](h T, j int) {
+	for {
+		i := (j - 1) / 2 // parent
+		if i == j || !h.Less(j, i) {
+			break
+		}
+		h.Swap(i, j)
+		j = i
+	}
+}
+
+func down[V any, T Interface[V]](h T, i0, n int) bool {
+	i := i0
+	for {
+		j1 := 2*i + 1
+		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
+			break
+		}
+		j := j1 // left child
+		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
+			j = j2 // = 2*i + 2  // right child
+		}
+		if !h.Less(j, i) {
+			break
+		}
+		h.Swap(i, j)
+		i = j
+	}
+	return i > i0
+}

kdtree.go

@@ -1,7 +1,6 @@
 package kdtree
 
 import (
-	"container/heap"
 	"fmt"
 	"sort"
 	"strings"
@@ -403,17 +402,18 @@ func nearestNeighbor[T Comparable[T]](d int, v, nn *T, cd int, r *kdNode[T]) *T
 	return nn
 }
 
-func (t *KDTree[T]) KNearestNeighbor(value T, k int) []T {
+func (t *KDTree[T]) KNearestNeighbor(value T, k int) []*T {
 	if t == nil || t.root == nil || t.size < k {
 		return nil
 	}
 
 	pqRes := NewBoundedPriorityQueue[T](k)
 	kNearestNeighbor(k, t.dimensions, &value, &pqRes, 0, t.root)
 
-	res := make([]T, 0, pqRes.Len())
-	for pqRes.Len() != 0 {
-		d := heap.Pop(&pqRes).(Item[T]).Data
+	res := make([]*T, 0, k)
+	for range k {
+		// heap with a preset capacity
+		d := internal.Pop(&pqRes).Data
 		res = append(res, d)
 	}
 
@@ -459,8 +459,8 @@ func kNearestNeighbor[T Comparable[T]](k, d int, v *T, pq *BoundedPriorityQueue[
 	ncd = (ncd - 1 + d) % d // Go back to the dimension used for splitting at the leaf node.
 	for path, cn, cDir := popLast(path); cn != nil; path, cn, cDir = popLast(path) {
 		currentDistance := (*v).Dist(cn.value)
-		heap.Push(pq, Item[T]{
-			Data:     cn.value,
+		internal.Push(pq, Item[T]{
+			Data:     &cn.value,
 			Priority: currentDistance,
 		})
 
@@ -478,7 +478,7 @@ func kNearestNeighbor[T Comparable[T]](k, d int, v *T, pq *BoundedPriorityQueue[
 }
 
 func getFarthestDistance[T Comparable[T]](pq *BoundedPriorityQueue[T]) int {
-	v := pq.Peek().(Item[T])
+	v := pq.Peek()
 	return v.Priority
 }
 

priorityQueue.go

@@ -1,12 +1,13 @@
 package kdtree
 
 import (
-	"container/heap"
 	"slices"
+
+	heap "github.com/rishitc/go-kd-tree/internal/utils"
 )
 
 type Item[T Comparable[T]] struct {
-	Data     T
+	Data     *T
 	Priority int
 }
 
@@ -17,7 +18,7 @@ type BoundedPriorityQueue[T Comparable[T]] struct {
 
 func NewBoundedPriorityQueue[T Comparable[T]](maxSize int) BoundedPriorityQueue[T] {
 	return BoundedPriorityQueue[T]{
-		data:     nil,
+		data:     make([]Item[T], 0, maxSize),
 		capacity: maxSize,
 	}
 }
@@ -33,8 +34,7 @@ func (pq BoundedPriorityQueue[T]) Swap(i, j int) {
 	pq.data[i], pq.data[j] = pq.data[j], pq.data[i]
 }
 
-func (pq *BoundedPriorityQueue[T]) Push(value any) {
-	item := value.(Item[T])
+func (pq *BoundedPriorityQueue[T]) Push(item Item[T]) {
 	isFull := pq.Len() == pq.capacity
 
 	if isFull {
@@ -46,14 +46,14 @@ func (pq *BoundedPriorityQueue[T]) Push(value any) {
 	pq.data = append(pq.data, item)
 }
 
-func (pq *BoundedPriorityQueue[T]) Pop() any {
+func (pq *BoundedPriorityQueue[T]) Pop() Item[T] {
 	n := pq.Len()
 	item := pq.data[n-1]
 	pq.data = slices.Delete(pq.data, n-1, n)
 	return item
 }
 
-func (pq *BoundedPriorityQueue[T]) Peek() any {
+func (pq *BoundedPriorityQueue[T]) Peek() Item[T] {
 	return pq.data[0]
 }
 

tests/kdtree_2d_test.go

@@ -184,7 +184,7 @@ func Test2DNearestNeighbor5(t *testing.T) {
 	}
 }
 
-func tensor2DSortFunc(a, b types.Tensor2D) int {
+func tensor2DSortFunc(a, b *types.Tensor2D) int {
 	if a[0] != b[0] {
 		return a[0] - b[0]
 	} else {
@@ -210,19 +210,19 @@ func Test2DKNearestNeighbor1(t *testing.T) {
 	tree := kdtree.NewKDTreeWithValues(dimensions, ps)
 	testTable := map[string]struct {
 		input    input
-		expected []types.Tensor2D
+		expected []*types.Tensor2D
 	}{
 		"Find the 2 closest neighbors to a point that is not in the KD tree.": {
 			input:    input{p: [2]int{25, 25}, k: 2},
-			expected: []types.Tensor2D{{40, 20}, {10, 25}},
+			expected: []*types.Tensor2D{{40, 20}, {10, 25}},
 		},
 		"The closest neighbor to a point that is in the KD tree.": {
 			input:    input{p: [2]int{60, 90}, k: 1},
-			expected: []types.Tensor2D{{60, 90}},
+			expected: []*types.Tensor2D{{60, 90}},
 		},
 		"The three closest neighbors to a point that is in the KD tree.": {
 			input:    input{p: [2]int{70, 70}, k: 3},
-			expected: []types.Tensor2D{{50, 50}, {60, 90}, {70, 70}},
+			expected: []*types.Tensor2D{{50, 50}, {60, 90}, {70, 70}},
 		},
 	}
 	for name, st := range testTable {

tests/priorityQueue_test.go

@@ -1,9 +1,10 @@
 package tests
 
 import (
-	"container/heap"
 	"testing"
 
+	heap "github.com/rishitc/go-kd-tree/internal/utils"
+
 	boundedpq "github.com/rishitc/go-kd-tree"
 	types "github.com/rishitc/go-kd-tree/internal/types"
 )
@@ -19,9 +20,9 @@ func TestBoundedPriorityQueue(t *testing.T) {
 		{
 			name: "More number of input elements as compared to the max allowed size of the bounded priority queue",
 			inputElems: []boundedpq.Item[types.Tensor2D]{
-				{Data: types.Tensor2D{1, 2}, Priority: 3},
-				{Data: types.Tensor2D{3, 5}, Priority: 2},
-				{Data: types.Tensor2D{6, 7}, Priority: 4},
+				{Data: &types.Tensor2D{1, 2}, Priority: 3},
+				{Data: &types.Tensor2D{3, 5}, Priority: 2},
+				{Data: &types.Tensor2D{6, 7}, Priority: 4},
 			},
 			inputMaxSize: 2,
 			expLen:       2,
@@ -33,9 +34,9 @@ func TestBoundedPriorityQueue(t *testing.T) {
 		{
 			name: "Less number of input elements as compared to the max allowed size of the bounded priority queue",
 			inputElems: []boundedpq.Item[types.Tensor2D]{
-				{Data: types.Tensor2D{1, 2}, Priority: 3},
-				{Data: types.Tensor2D{3, 5}, Priority: 2},
-				{Data: types.Tensor2D{6, 7}, Priority: 4},
+				{Data: &types.Tensor2D{1, 2}, Priority: 3},
+				{Data: &types.Tensor2D{3, 5}, Priority: 2},
+				{Data: &types.Tensor2D{6, 7}, Priority: 4},
 			},
 			inputMaxSize: 5,
 			expLen:       3,
@@ -48,9 +49,9 @@ func TestBoundedPriorityQueue(t *testing.T) {
 		{
 			name: "Equal number of input elements as compared to the max allowed size of the bounded priority queue",
 			inputElems: []boundedpq.Item[types.Tensor2D]{
-				{Data: types.Tensor2D{1, 2}, Priority: 3},
-				{Data: types.Tensor2D{3, 5}, Priority: 2},
-				{Data: types.Tensor2D{6, 7}, Priority: 4},
+				{Data: &types.Tensor2D{1, 2}, Priority: 3},
+				{Data: &types.Tensor2D{3, 5}, Priority: 2},
+				{Data: &types.Tensor2D{6, 7}, Priority: 4},
 			},
 			inputMaxSize: 3,
 			expLen:       3,
@@ -63,10 +64,10 @@ func TestBoundedPriorityQueue(t *testing.T) {
 		{
 			name: "Inserting two elements with the same priority and the one inserted first should be retained",
 			inputElems: []boundedpq.Item[types.Tensor2D]{
-				{Data: types.Tensor2D{1, 2}, Priority: 2},
-				{Data: types.Tensor2D{3, 5}, Priority: 3},
-				{Data: types.Tensor2D{4, 8}, Priority: 4},
-				{Data: types.Tensor2D{6, 7}, Priority: 4},
+				{Data: &types.Tensor2D{1, 2}, Priority: 2},
+				{Data: &types.Tensor2D{3, 5}, Priority: 3},
+				{Data: &types.Tensor2D{4, 8}, Priority: 4},
+				{Data: &types.Tensor2D{6, 7}, Priority: 4},
 			},
 			inputMaxSize: 3,
 			expLen:       3,
@@ -94,7 +95,7 @@ func TestBoundedPriorityQueue(t *testing.T) {
 			}
 
 			for i := range st.expLen {
-				headElem := pq.Peek().(boundedpq.Item[types.Tensor2D]).Data
+				headElem := pq.Peek().Data
 				expElem := st.expOrderedElems[i]
 				if expElem[0] != headElem[0] && expElem[1] != headElem[1] {
 					t.Errorf("Expected next element to be %v, got %v", expElem, headElem)