在Java中，如何查看（）PriorityQueue中的前k个元素？

Question

I see in the documentation, that PriorityQueue.peek() gives me access to the head of the queue in O(1), but what if I need access to the k top elements in the queue ? I would use poll() k times, but that takes O(log(N)) , is there a way to do it in constant time ?

Answer 1

Nope. If you could do it in constant time, you could do a comparison sort in linear time by heapifying an array and then finding the top N items, where N is all of them.

Answer 2

It's not possible to do it in constant time. If your data is already in a PriorityQueue , removing the top k elements one by one is the best you can do. Each remove costs you O(log(n)) but you have also figured that one out, hence the question.

However, if you are not forced to use a PriorityQueue , then you can do a partial sort on a list and retrieve the top k elements. The asymptotic complexity of a partial sort is O(n*log(k)) . It can be executed faster than the PriorityQueue approach if we also take into account the cost of setting up the priority queue, see Selecting top k items from a list efficiently in Java / Groovy (select the top 5 elements from a list of 10 million, PriorityQueue 300ms vs. partial sort 170ms).

Answer 3

It can be done in O(klogk) time and O(k) additional space if the input is represented as an indexed sequence in heap order. If k << n , that's a very good improvement.

static List <Integer> kLargest(List<Integer> xs, int k) {
 if (k <= 0) return Collections.emptyList();

 PriorityQueue<Map.Entry<Integer, Integer>> pq = 
    new PriorityQueue<>(Comparator.comparingInt(Map.Entry::getValue));

 pq.add(new SimpleImmutableEntry<>(0, xs.get(0)));
 List<Integer> result = new ArrayList<>();
 for (int i = 0; i < k; i++) {
  Map.Entry <Integer, Integer> max = pq.remove();
  result.add(max.getValue());

  int j = max.getKey();
  int left = 2 * j + 1;
  int right = 2 * j + 2;

  if (left < xs.size()) pq.add(new SimpleImmutableEntry<>(left, xs.get(left)));
  if (right < xs.size()) pq.add(new SimpleImmutableEntry<>(right, xs.get(right)));
 }
 return result;
}

The question is, how to get the input in heap order. If you create your own PriorityQueue class, you'll obviously have access to the internal array.