K Smallest In Unsorted Array

Find the K smallest numbers in an unsorted integer array A. The returned numbers should be in ascending order.

Assumptions

• A is not null
• K is >= 0 and smaller than or equal to size of A

Return

• an array with size K containing the K smallest numbers in ascending order

Examples

• A = {3, 4, 1, 2, 5}, K = 3, the 3 smallest numbers are {1, 2, 3}

public class Solution {
  // use a priority queue(max heap with size of k) to figure it out.
  // we traverse the array one time, whenever we met a element, we check it if smaller then the top element of the heap
  // if it is, we poll and insert that new element in
  // if it isn't, we do nothing
  // time: O(n)
  // space: O(k)
  //
  // thing to konw better: the vanilla priority q in java is max heap or min heap?
  // how to turn a priority q into a array or a list in a efficient way?
  //
  // assumption: all element in array is integer and small than Integer.MAX_VALUE and larger than Integer.MIN_VALUE
  // array is not null, array's length can be 0
  // k can be zero, k is not larger than the length of array
  public int[] kSmallest(int[] array, int k) {
    // Write your solution here
    int[] res = new int[k];
    if(k==0){
      return res;
    }
    PriorityQueue<Integer> pq = new PriorityQueue<>(k,Collections.reverseOrder());
    for(int i=0; i<array.length; i++){
      if(pq.size()<k){
        pq.offer(array[i]);
      }else{
        if(array[i]<pq.peek()){
          pq.poll();
          pq.offer(array[i]);
        }
      }
    }
    for(int i=0; i<k; i++){
      res[k-i-1] = pq.poll();
    }
    return res;
  }
}