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;
}
}