Question 3. Finding smallest k numbers

阅读更多

Question: Among n integers , to find k smallest ones.

Sample Input : 5, 2, 1, 3 , 4, 7, 8, 6

Sample output : the 3 smallest ones are : 1, 2, 3.

 

 

My thought: Maximum-Heap is the first data structure comes to me. We can construct a maximum heap for the first k numbers and for each next number we compare it with the number on top of the heap, if it's bigger, we ignore it, otherwise we insert it into the heap and remove the number on top of the heap. Finally the numbers left in the heap is the smallest k numbers. The time complexity should be O(n log k)

 

Java Implementation (Java has its own implementation of heap as PriorityQueue. While here , for practice, we just provide a very specific maximum integer heap.): 

 

public class KSmallest {

	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) {
		MaxIntHeap heap = new MaxIntHeap(nums, 0, k);
		for ( int i = k ; i < nums.length ; i ++ ) {
			if (nums[i] < heap.peek()) {
				heap.pop();
				heap.add(nums[i]);
			}
		}
		
		return heap.getAll();
	}
	
	
	
	public static class MaxIntHeap {
		
		private int[] nums;
		private int count;

		
		public MaxIntHeap ( int[] nums, int offset, int count){
		   this.nums = new int[count+1]; // index 0 will not be used
		   System.arraycopy(nums, offset, this.nums, 1, count);
		   this.count = count;
		   for ( int i = count/2 ; i > 0 ; i --) {
			   swimDown(i);
		   }
		}
		
		public int[] getAll() {
			int result[] = new int[count];
			System.arraycopy(this.nums, 1, result, 0, count);
			return result;
		}
	    
		
		private void swimUp (int i) {
			int parent = i/2;
			while (parent != 0 && nums[i] > nums[parent]) {
				swap ( i, parent);
				i = parent;
				parent = i/2;
			}
		}
		
		private void swimDown(int i) {
			int lchild = i * 2;
			if ( lchild > count )
				return;
			int bigger = lchild;
			int rchild = lchild + 1;
			if (rchild <= count && nums[rchild] > nums[lchild]) 
				bigger = rchild; 
			while( nums[bigger] > nums[i]) {
				swap ( bigger, i);
				i = bigger;
				lchild = i * 2;
				if ( lchild > count )
					return;
				bigger = lchild;
				rchild = lchild + 1;
				if (rchild <= count && nums[rchild] > nums[lchild]) 
					bigger = rchild; 
			}
			
		}
		
		private void swap (int i , int j) {
			int temp = nums[i];
			nums[i] = nums[j];
			nums[j] = temp;
		}
		
		public int peek() {
			return nums[1];
		}
		
		public int pop() {
			swap(1, count--);
			swimDown(1);
			return nums[count + 1];
		}
		
		//don't concern resizing array here
		public void add (int num) {
			nums[++count] = num;
			swimUp(count);
		}
	}

}