fix(#15309): Optimize PriorityQueue.remove() function

Author: vijaykriishna

lucene/core/src/java/org/apache/lucene/util/PriorityQueue.java

@@ -19,7 +19,9 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.Iterator;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.function.IntFunction;
 import java.util.function.Supplier;
@@ -72,6 +74,7 @@ public static <T> PriorityQueue<T> usingComparator(
   private final int maxSize;
   private final T[] heap;
   private final LessThan<? super T> lessThan;
+  private final Map<T, Integer> indexMap = new HashMap<>();
 
   /** Create an empty priority queue of the configured size using the specified {@link LessThan}. */
   public PriorityQueue(int maxSize, LessThan<? super T> lessThan) {
@@ -272,6 +275,7 @@ public final int size() {
   public final void clear() {
     Arrays.fill(heap, 0, size + 1, null);
     size = 0;
+    indexMap.clear();
   }
 
   /**
@@ -280,20 +284,21 @@ public final void clear() {
    * constant remove time but the trade-off would be extra cost to all additions/insertions)
    */
   public final boolean remove(T element) {
-    for (int i = 1; i <= size; i++) {
-      if (heap[i] == element) {
-        heap[i] = heap[size];
-        heap[size] = null; // permit GC of objects
-        size--;
-        if (i <= size) {
-          if (!upHeap(i)) {
-            downHeap(i);
-          }
-        }
-        return true;
+    Integer idx = indexMap.get(element);
+    if (idx == null || idx > size || heap[idx] != element) return false;
+    heap[idx] = heap[size];
+    indexMap.remove(element);
+    if (idx != size) {
+      indexMap.put(heap[idx], idx);
+    }
+    heap[size] = null;
+    size--;
+    if (idx <= size) {
+      if (!upHeap(idx)) {
+        downHeap(idx);
       }
     }
-    return false;
+    return true;
   }
 
   /**
@@ -320,36 +325,43 @@ public T[] drainToArrayHighestFirst(IntFunction<T[]> newArray) {
     return array;
   }
 
-  private boolean upHeap(int origPos) {
-    int i = origPos;
-    T node = heap[i]; // save bottom node
-    int j = i >>> 1;
-    while (j > 0 && lessThan.lessThan(node, heap[j])) {
-      heap[i] = heap[j]; // shift parents down
-      i = j;
-      j = j >>> 1;
-    }
-    heap[i] = node; // install saved node
-    return i != origPos;
+  protected final void addToHeap(int idx, T element) {
+    heap[idx] = element;
+    indexMap.put(element, idx);
   }
 
-  private void downHeap(int i) {
-    T node = heap[i]; // save top node
-    int j = i << 1; // find smaller child
-    int k = j + 1;
-    if (k <= size && lessThan.lessThan(heap[k], heap[j])) {
-      j = k;
+  protected boolean upHeap(int i) {
+    T node = heap[i];
+    int j = i;
+    while (j > 1 && lessThan.lessThan(node, heap[j >> 1])) {
+      heap[j] = heap[j >> 1];
+      indexMap.put(heap[j], j);
+      j >>= 1;
     }
-    while (j <= size && lessThan.lessThan(heap[j], node)) {
-      heap[i] = heap[j]; // shift up child
-      i = j;
-      j = i << 1;
-      k = j + 1;
-      if (k <= size && lessThan.lessThan(heap[k], heap[j])) {
+    heap[j] = node;
+    indexMap.put(node, j);
+    return j < i;
+  }
+
+  protected boolean downHeap(int i) {
+    T node = heap[i];
+    int j = i;
+    int k;
+    while ((k = j << 1) <= size) {
+      if (k < size && lessThan.lessThan(heap[k + 1], heap[k])) {
+        k++;
+      }
+      if (lessThan.lessThan(heap[k], node)) {
+        heap[j] = heap[k];
+        indexMap.put(heap[j], j);
         j = k;
+      } else {
+        break;
       }
     }
-    heap[i] = node; // install saved node
+    heap[j] = node;
+    indexMap.put(node, j);
+    return j > i;
   }
 
   /**