Top “n” using a Priority Queue

Biju KunjummenMarch 3rd, 2021Last Updated: March 1st, 2021

0 322 2 minutes read

If you ever need to capture the smallest or largest “n” from a stream of data, the approach more often than not will be to use a simple data-structure called the Priority Queue.

Priority Queues do one thing very well – once a bunch of data is added, it can return the lowest value (or the highest value) in constant time.

How is this useful to answer a top or bottom “n” type question. Let’s see.

Consider this hypothetical stream of data:

And you have to answer the smallest 2 at any point as this stream of data comes in. The trick is to use a Priority Queue

with a size limit of 2
which returns the largest of these 2 when asked for (sometimes referred to as a Max Priority Queue)

Two considerations as data streams in:

if the size of the priority queue is less than 2 then add the value to the priority queue
if the size of the priority queue is equal to 2, then compare the value from the stream with the largest in the queue
- if less then remove the largest and add the new value
- if more then do nothing

At the bottom is the state of the Priority Queue as each data in the stream is processed:

See how it always holds the bottom 2.

Similarly for the largest 3, a Priority Queue with a max capacity of 3 which returns the smallest (referred to as a Min Priority Queue) can be used the following way:

if the size of the priority queue is less than 3, then add to the priority queue
if the size is equal to 2, then check the value from the stream with the smallest in the queue
- if more then remove smallest add the value from stream and ignore otherwise

Implementation

Here is a simple kotlin based implementation that uses the built in
PriorityQueue in Java standard library.

fun findNSmallestAndLargest(nums: List<Int>, n: Int): Pair<List<Int>, List<Int>> {
    val minFirst: Comparator<Int> = Comparator.naturalOrder<Int>()
    val maxFirst: Comparator<Int> = minFirst.reversed()
    val minPq: PriorityQueue<Int> = PriorityQueue(minFirst)
    val maxPq: PriorityQueue<Int> = PriorityQueue(maxFirst)
    for (num in nums) {
        checkAndAddIfSmallest(maxPq, n, num)
        checkAndAddIfLargest(minPq, n, num)
    }
    return maxPq.toList() to minPq.toList()
}
 
private fun checkAndAddIfSmallest(maxPq: PriorityQueue<Int>, n: Int, num: Int) {
    if (maxPq.size < n) {
        maxPq.add(n)
    } else if (num < maxPq.peek()) {
        maxPq.poll()
        maxPq.add(num)
    }
}
 
private fun checkAndAddIfLargest(minPq: PriorityQueue<Int>, n: Int, num: Int) {
    if (minPq.size < n) {
        minPq.add(n)
    } else if (num > minPq.peek()) {
        minPq.poll()
        minPq.add(num)
    }
}

The implementation is very straightforward and follows the outlined algorithm to the letter.

Published on Java Code Geeks with permission by Biju Kunjummen, partner at our JCG program. See the original article here: Top “n” using a Priority Queue

Opinions expressed by Java Code Geeks contributors are their own.