return num of neighbors

keilerkonzept · Oct 10, 2024 · d1f33b3 · d1f33b3
1 parent 4b9460d
commit d1f33b3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/lsh/model_wide.go b/lsh/model_wide.go
@@ -84,7 +84,7 @@ func (me *WideModel) Predict1Alloc(k int, x []uint64, votes bitknn.Votes) int {
 // Predict1Into predicts the label for a single input using the given slices (of length [k]+1 each) for the neighbor heaps.
 func (me *WideModel) Predict1Into(k int, x []uint64, votes bitknn.Votes, bucketDistances []int, bucketIDs []uint64, distances []int, indices []int) int {
 	xp := me.Hash.Hash1Wide(x)
-	k, n := NearestWide(me.WideData, me.BucketIDs, me.Buckets, k, xp, x, bucketDistances, bucketIDs, distances, indices)
-	me.WideModel.Narrow.Vote(k, distances, indices, votes)
-	return n
+	k0, _ := NearestWide(me.WideData, me.BucketIDs, me.Buckets, k, xp, x, bucketDistances, bucketIDs, distances, indices)
+	me.WideModel.Narrow.Vote(k0, distances, indices, votes)
+	return k0
 }
diff --git a/model_wide.go b/model_wide.go
@@ -22,13 +22,16 @@ func (me *WideModel) PreallocateHeap(k int) {
 }
 
 // Predicts the label of a single input point. Reuses two slices of length K+1 for the neighbor heap.
-func (me *WideModel) Predict1(k int, x []uint64, votes Votes) {
+// Returns the number of neighbors found.
+func (me *WideModel) Predict1(k int, x []uint64, votes Votes) int {
 	me.Narrow.PreallocateHeap(k)
-	me.Predict1Into(k, x, me.Narrow.HeapDistances, me.Narrow.HeapIndices, votes)
+	return me.Predict1Into(k, x, me.Narrow.HeapDistances, me.Narrow.HeapIndices, votes)
 }
 
 // Predicts the label of a single input point, using the given slices for the neighbor heap.
-func (me *WideModel) Predict1Into(k int, x []uint64, distances []int, indices []int, votes Votes) {
+// Returns the number of neighbors found.
+func (me *WideModel) Predict1Into(k int, x []uint64, distances []int, indices []int, votes Votes) int {
 	k = NearestWide(me.WideData, k, x, distances, indices)
 	me.Narrow.Vote(k, distances, indices, votes)
+	return k
 }