Tuesday 22 September 2015

Maximum k integers using min heap

Problem

Find k maximum integers from an array of infinite integers.

Brute force

We can sort the integers and return the top k elements. The complexity is O(n log n) because of sorting. As n is very large we would have to store all the elements in memory for sorting them.

Better solution

We can implement a min heap to hold exactly k elements. Every time a new number is added we can check with the top of the element which is the minimum of the current k elements. If new number is less than the min heap top then that number is discarded otherwise it is added in the heap after removing the current top. In this way we always keep only k items in memory and the complexity is O(n log k) (assumption is that k is significantly lesser than n)

Code

import java.util.PriorityQueue;

public class MaxKUsingMinHeap
{
 public static void main(String[] args)
 {
  int[] arr =
  { 3, 46, 2, 56, 3, 38, 93, 45, 6, 787, 34, 76, 
    44, 6, 7, 86, 8, 44, 56 };
  int[] result = getTopElements(arr, 5);
  for (int i : result)
  {
   System.out.print(i + ",");
  }
 }

 public static int[] getTopElements(int[] arr, int k)
 {
  PriorityQueue<Integer> minHeap = 
   new PriorityQueue<Integer>();
  for (int i = 0; i < arr.length; ++i)
  {
   int currentNum = arr[i];
   if (minHeap.size() < k)
    minHeap.add(currentNum);
   else if (currentNum > minHeap.peek())
   {
    minHeap.poll();
    minHeap.add(currentNum);
   }
  }
  int[] result = new int[minHeap.size()];
  int index = 0;
  while (!minHeap.isEmpty())
  {
   result[index++] = minHeap.poll();
  }
  return result;
 }
}
        

No comments:

Post a Comment