feat: ✨ add support for finding the k nearest neighbors

Author: rishitc

go.mod

@@ -1,6 +1,6 @@
 module github.com/rishitc/go-kd-tree
 
-go 1.21.4
+go 1.22.5
 
 require (
 	github.com/google/flatbuffers v24.3.25+incompatible

kdtree.go

@@ -1,6 +1,7 @@
 package kdtree
 
 import (
+	"container/heap"
 	"fmt"
 	"sort"
 	"strings"
@@ -408,6 +409,93 @@ 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 {
+	if t == nil || t.root == nil || t.sz < 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 = append(res, d)
+	}
+
+	return res
+}
+
+type direction bool
+
+const (
+	left  direction = true
+	right           = false
+)
+
+type nodeInfo[T Comparable[T]] struct {
+	node *kdNode[T]
+	dir  direction
+}
+
+func kNearestNeighbor[T Comparable[T]](k, d int, v *T, pq *BoundedPriorityQueue[T], cd int, r *kdNode[T]) {
+	if r == nil {
+		return
+	}
+
+	ncd := cd
+
+	var path []nodeInfo[T]
+	for r != nil {
+		info := nodeInfo[T]{
+			node: r,
+		}
+		if rel := (*v).Order(r.value, ncd); rel == Lesser {
+			r = r.left
+			info.dir = left
+		} else {
+			r = r.right
+			info.dir = right
+		}
+		path = append(path, info)
+
+		ncd = (ncd + 1) % d
+	}
+
+	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,
+			Priority: currentDistance,
+		})
+
+		if pq.Len() < pq.Capacity() || (*v).DistDim(cn.value, ncd) < getFarthestDistance(pq) {
+			var next *kdNode[T]
+			if cDir == left {
+				next = cn.right
+			} else {
+				next = cn.left
+			}
+			kNearestNeighbor(k, d, v, pq, (ncd+1)%d, next)
+		}
+		ncd = (ncd - 1 + d) % d
+	}
+}
+
+func getFarthestDistance[T Comparable[T]](pq *BoundedPriorityQueue[T]) int {
+	v := pq.Peek().(Item[T])
+	return v.Priority
+}
+
+func popLast[T Comparable[T]](arr []nodeInfo[T]) ([]nodeInfo[T], *kdNode[T], direction) {
+	if len(arr) == 0 {
+		return arr, nil, left
+	}
+	li := len(arr) - 1
+	return arr[:li], arr[li].node, arr[li].dir
+}
+
 func findMin[T Comparable[T]](d, tcd, cd int, r *kdNode[T]) *T {
 	if r == nil {
 		return nil

kdtree_2d_test.go

@@ -251,6 +251,61 @@ func Test2DNearestNeighbor5(t *testing.T) {
 	}
 }
 
+func tensor2DSortFunc(a, b Tensor2D) int {
+	if a[0] != b[0] {
+		return a[0] - b[0]
+	} else {
+		return a[1] - b[1]
+	}
+}
+
+func Test2DKNearestNeighbor1(t *testing.T) {
+	const dimensions = 2
+	type input struct {
+		p Tensor2D
+		k int
+	}
+	ps := []Tensor2D{
+		{50, 50},
+		{10, 25},
+		{40, 20},
+		{25, 80},
+		{70, 70},
+		{60, 10},
+		{60, 90},
+	}
+	tree := kdtree.NewKDTreeWithValues(dimensions, ps)
+	testTable := map[string]struct {
+		input    input
+		expected []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: []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: []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: []Tensor2D{{50, 50}, {60, 90}, {70, 70}},
+		},
+	}
+	for name, st := range testTable {
+		t.Run(name, func(t *testing.T) {
+			nns := tree.KNearestNeighbor(st.input.p, st.input.k)
+			slices.SortFunc(nns, tensor2DSortFunc)
+			slices.SortFunc(st.expected, tensor2DSortFunc)
+			for i := range len(st.expected) {
+				if st.expected[i][0] != nns[i][0] || st.expected[i][1] != nns[i][1] {
+					t.Fatalf("Expected point: %v, got %v", st.expected[i], nns[i])
+				}
+			}
+		})
+	}
+}
+
 func Test2DNodeAddition1(t *testing.T) {
 	const dimensions = 2
 	ps := []Tensor2D{

priorityQueue.go

@@ -0,0 +1,62 @@
+package kdtree
+
+import (
+	"container/heap"
+	"slices"
+)
+
+type Item[T Comparable[T]] struct {
+	Data     T
+	Priority int
+}
+
+type BoundedPriorityQueue[T Comparable[T]] struct {
+	data     []Item[T]
+	capacity int
+}
+
+func NewBoundedPriorityQueue[T Comparable[T]](maxSize int) BoundedPriorityQueue[T] {
+	return BoundedPriorityQueue[T]{
+		data:     nil,
+		capacity: maxSize,
+	}
+}
+
+func (pq BoundedPriorityQueue[T]) Len() int { return len(pq.data) }
+
+func (pq BoundedPriorityQueue[T]) Less(i, j int) bool {
+	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
+	return pq.data[i].Priority > pq.data[j].Priority
+}
+
+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])
+	isFull := pq.Len() == pq.capacity
+
+	if isFull {
+		if pq.data[0].Priority <= item.Priority {
+			return
+		}
+		heap.Pop(pq)
+	}
+	pq.data = append(pq.data, item)
+}
+
+func (pq *BoundedPriorityQueue[T]) Pop() any {
+	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 {
+	return pq.data[0]
+}
+
+func (pq *BoundedPriorityQueue[T]) Capacity() int {
+	return pq.capacity
+}

priorityQueue_test.go

@@ -0,0 +1,108 @@
+package kdtree_test
+
+import (
+	"container/heap"
+	"testing"
+
+	kdtree "github.com/rishitc/go-kd-tree"
+)
+
+// This example creates a PriorityQueue with some items, adds and manipulates an item,
+// and then removes the items in priority order.
+func TestBoundedPriorityQueue(t *testing.T) {
+	table := []struct {
+		name            string
+		inputElems      []kdtree.Item[Tensor2D] // Some items and their priorities.
+		inputMaxSize    int
+		expLen          int
+		expOrderedElems []Tensor2D
+	}{
+		{
+			name: "More number of input elements as compared to the max allowed size of the bounded priority queue",
+			inputElems: []kdtree.Item[Tensor2D]{
+				{Data: Tensor2D{1, 2}, Priority: 3},
+				{Data: Tensor2D{3, 5}, Priority: 2},
+				{Data: Tensor2D{6, 7}, Priority: 4},
+			},
+			inputMaxSize: 2,
+			expLen:       2,
+			expOrderedElems: []Tensor2D{
+				{1, 2},
+				{3, 5},
+			},
+		},
+		{
+			name: "Less number of input elements as compared to the max allowed size of the bounded priority queue",
+			inputElems: []kdtree.Item[Tensor2D]{
+				{Data: Tensor2D{1, 2}, Priority: 3},
+				{Data: Tensor2D{3, 5}, Priority: 2},
+				{Data: Tensor2D{6, 7}, Priority: 4},
+			},
+			inputMaxSize: 5,
+			expLen:       3,
+			expOrderedElems: []Tensor2D{
+				{6, 7},
+				{1, 2},
+				{3, 5},
+			},
+		},
+		{
+			name: "Equal number of input elements as compared to the max allowed size of the bounded priority queue",
+			inputElems: []kdtree.Item[Tensor2D]{
+				{Data: Tensor2D{1, 2}, Priority: 3},
+				{Data: Tensor2D{3, 5}, Priority: 2},
+				{Data: Tensor2D{6, 7}, Priority: 4},
+			},
+			inputMaxSize: 3,
+			expLen:       3,
+			expOrderedElems: []Tensor2D{
+				{6, 7},
+				{1, 2},
+				{3, 5},
+			},
+		},
+		{
+			name: "Inserting two elements with the same priority and the one inserted first should be retained",
+			inputElems: []kdtree.Item[Tensor2D]{
+				{Data: Tensor2D{1, 2}, Priority: 2},
+				{Data: Tensor2D{3, 5}, Priority: 3},
+				{Data: Tensor2D{4, 8}, Priority: 4},
+				{Data: Tensor2D{6, 7}, Priority: 4},
+			},
+			inputMaxSize: 3,
+			expLen:       3,
+			expOrderedElems: []Tensor2D{
+				{4, 8},
+				{3, 5},
+				{1, 2},
+			},
+		},
+	}
+
+	for _, st := range table {
+		t.Run(st.name, func(t *testing.T) {
+			pq := kdtree.NewBoundedPriorityQueue[Tensor2D](st.inputMaxSize)
+			for _, elem := range st.inputElems {
+				data, priority := elem.Data, elem.Priority
+				heap.Push(&pq, kdtree.Item[Tensor2D]{
+					Data:     data,
+					Priority: priority,
+				})
+			}
+
+			if pq.Len() != st.expLen {
+				t.Errorf("Expected size to be %v, got %v", st.expLen, pq.Len())
+			}
+
+			for i := range st.expLen {
+				headElem := pq.Peek().(kdtree.Item[Tensor2D]).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)
+				}
+				heap.Pop(&pq)
+			}
+		})
+	}
+}
+