From 04604f6b302a7fc93f2fad85651969e9a8586392 Mon Sep 17 00:00:00 2001
From: Runqi Yang <runqiyang@gmail.com>
Date: Thu, 11 Oct 2018 14:15:38 +0800
Subject: [PATCH] support search_k (issue #3)

Implements support for search_k. Keeps the original API and make the default value of searchK -1 (same as C++ and python API).
---
 src/main/java/com/spotify/annoy/ANNIndex.java | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/spotify/annoy/ANNIndex.java b/src/main/java/com/spotify/annoy/ANNIndex.java
index 20413ec..01fb1bf 100644
--- a/src/main/java/com/spotify/annoy/ANNIndex.java
+++ b/src/main/java/com/spotify/annoy/ANNIndex.java
@@ -236,6 +236,11 @@ private static boolean isZeroVec(float[] v) {
   @Override
   public final List<Integer> getNearest(final float[] queryVector,
                                         final int nResults) {
+    return getNearest(queryVector, nResults, -1);
+  }
+
+  public final List<Integer> getNearest(final float[] queryVector,
+                                        final int nResults, int searchK) {
 
     if (queryVector.length != DIMENSION) {
       throw new RuntimeException(String.format("queryVector must be size of %d, but was %d",
@@ -250,8 +255,12 @@ public final List<Integer> getNearest(final float[] queryVector,
       pq.add(new PQEntry(kMaxPriority, r));
     }
 
+    if (searchK == -1) {
+      searchK = roots.size() * nResults;
+    }
+
     Set<Integer> nearestNeighbors = new HashSet<Integer>();
-    while (nearestNeighbors.size() < roots.size() * nResults && !pq.isEmpty()) {
+    while (nearestNeighbors.size() < searchK && !pq.isEmpty()) {
       PQEntry top = pq.poll();
       long topNodeOffset = top.nodeOffset;
       int nDescendants = getIntInAnnBuf(topNodeOffset);