Merge pull request #19 from LeagueCode/nearset-pqEntries

Update get Nearset pq entries

Author: Erik Bernhardsson

src/main/java/com/spotify/annoy/ANNIndex.java

@@ -228,8 +228,17 @@ public class PQEntry implements Comparable<PQEntry> {
  this.nodeOffset = nodeOffset;
     }
 
- private float margin;
- private long nodeOffset;
+ private final float margin;
+ private final long nodeOffset;
+
+ public float getMargin() {
+ return margin;
+    }
+
+ public long getNodeOffset() {
+ return nodeOffset;
+    }
+
 
  @Override
  public int compareTo(final PQEntry o) {
@@ -254,17 +263,21 @@ public final List<Integer> getNearest(final float[] queryVector,
  public final List<Integer> getNearest(final float[] queryVector,
  final int nResults, int searchK) {
 
- ArrayList<PQEntry> sortedNNs = getNearestPqEntries(queryVector, nResults);
+ List<PQEntry> resultingPQEntries = getNearestPqEntries(queryVector, nResults, searchK);
 
  ArrayList<Integer> result = new ArrayList<>(nResults);
- for (int i = 0; i < nResults && i < sortedNNs.size(); i++) {
- result.add((int) sortedNNs.get(i).nodeOffset);
+ for (PQEntry pqEntry: resultingPQEntries) {
+ result.add((int) pqEntry.nodeOffset);
     }
  return result;
   }
 
- @Override
- public ArrayList<PQEntry> getNearestPqEntries(final float[] queryVector, final int nResults ) {
+ @Override
+ public final List<PQEntry> getNearestPqEntries( final float[] queryVector, final int nResults ) {
+ return getNearestPqEntries(queryVector, nResults, -1);
+   }
+ 
+ public List<PQEntry> getNearestPqEntries(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",
  DIMENSION, queryVector.length));
@@ -325,7 +338,8 @@ public ArrayList<PQEntry> getNearestPqEntries(final float[] queryVector, final i
       }
     }
  Collections.sort(sortedNNs);
- return sortedNNs;
+
+ return sortedNNs.subList(0, nResults);
   }
 
  /**