lsh: only stop early for 1-NN

keilerkonzept · Oct 10, 2024 · ed9b0db · ed9b0db
1 parent e728441
commit ed9b0db
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/lsh/nearest.go b/lsh/nearest.go
@@ -27,8 +27,8 @@ func Nearest(data []uint64, bucketIDs []uint64, buckets map[uint64]slice.IndexRa
 	numExamined := exactBucket.Length
 	nearestInBucket(data, exactBucket, k, x, &distances[0], &dataHeap)
 
-	// if the exact bucket already contains k neighbors, stop and return them
-	if dataHeap.Len() == k {
+	// stop early for 1-NN
+	if k == 1 && dataHeap.Len() == k {
 		return k, exactBucket.Length
 	}
 

diff --git a/lsh/nearest_wide.go b/lsh/nearest_wide.go
@@ -13,6 +13,7 @@ func NearestWide(data [][]uint64, bucketIDs []uint64, buckets map[uint64]slice.I
 	exactBucket := buckets[xh]
 	numExamined := exactBucket.Length
 	nearestWideInBucket(data, exactBucket, k, x, &distances[0], &dataHeap)
+
 	// stop early for 1-NN
 	if k == 1 && dataHeap.Len() == k {
 		return k, exactBucket.Length