Question 3b Smallest K numbers

Continue reading the solutions , I found a very clever optimization to the previous algorithm: 

 

After heapfying the whole array into a minimum heap and when poping up minimum elements from the heap , it's not necessary to sink down the last element ( the last element will be swapped with the first element in the pop operation) to the bottom , you just need to sink down it for k-1 times ( and then in the next pop k-2 times and so on), Because we just need to find the k smallest elements , for those elements under the layer k they are unrelated.  (i.e. After first pop operation, the last element will only sink k-1 times and the whole heap is not a minimum heap anymore but the k-1 smallest elements are still guaranteed to be in level 1 - level k-1 which the sub heap is still a minimum heap.)

 

It's a very specific optimization to this scenario and the time complexity is O(2n + k^2)  ( for the pop operation the time spent is 1 + 2 + 3 + ... + k-1 ).

 

Java Implementation:

public class KSmallestB {

	public static void main(String[] args) {
		 assertEquals(new int[]{1,2,3,4}, getKSmallest(new int[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, 4));
         assertEquals(new int[]{1,2,3,4}, getKSmallest(new int[]{10, 9, 8, 7, 6, 5, 4, 3, 2, 1}, 4));
         assertEquals(new int[]{21,22,14,18,9}, getKSmallest(new int[]{27, 14, 18, 22, 21, 91, 33, 36, 42, 78 , 9, 65, 101, 29}, 5));
	}
	
	private static void assertEquals(int[] standard, int[] result) {
		Arrays.sort(standard);
		Arrays.sort(result);
		assert Arrays.equals(standard, result);
		
	}
	
	public static int[] getKSmallest(int[] nums , int k) {
		
		heapfy (nums);
		int[] result = new int[k];
		for (int i = k-1 ; i >= 0 ; i --) {
			// nums.length + i - (k-1) is the total count left in the heap
			result[i] = pop(nums, nums.length + i - (k-1)  , i);
		}
		
		return result;
	}
	
	public static void heapfy (int[] nums) {
		for (int i = nums.length /2 ; i > 0 ; i -- ) {
			sink(nums, nums.length, i, Integer.MAX_VALUE);
		}
	}
	
	// return the smallest and sink at most depth times for the last element
	public static int pop ( int[] nums, int length, int depth) {
		swap(nums, 1, length--);
		sink(nums, length, 1 , depth);
		return nums[length];
	}
	
	// sink at most depth times
	public static void sink ( int[] nums , int length, int i , int depth) {
		int j = i * 2;
		while( j <= length && depth> 0 ) {
			// -1 is to adjust the index , because this array starts from 0
			if ( j < length && nums[j+1 -1] < nums[j -1]) j++;
			if (nums[j -1] < nums[i -1]) {
				swap(nums, i, j);
				i = j;
				j = 2*i;
			} else {
				break;
			}
		}
		
	}
	
	public static void swap ( int[] nums, int i , int j) {
		int temp = nums[i-1];
		nums[i-1] = nums[j-1];
		nums[j-1] = temp;
	}
	
	

}