The smallest K number

The smallest K number

Question link

Niuke.com

Problem-solving ideas

Minimum heap of size K

• Complexity: O(NlogK) + O(K)
• Particularly suitable for processing large amounts of data

The process of maintaining a minimum heap of size K is as follows: Use a big top heap. After adding an element, if the size of the big top heap is greater than K, then the top element of the big top heap is removed, that is, the element with the largest value in the current heap is removed, so that the elements remaining in the heap are smaller than those removed. The elements come in small sizes.

A large top heap should be used to maintain a minimum heap, rather than directly creating a small top heap and setting a size in an attempt to make the elements in the small top heap all minimum elements.

Java's PriorityQueue implements heap capabilities. PriorityQueue defaults to a small top heap. You can use the Lambda expression (o1, o2) -> o2 - o1 during initialization to implement a large top heap. Other languages ​​have similar heap data structures.

public ArrayList<Integer> GetLeastNumbers_Solution(int[] nums, int k) {
    
    
    if (k > nums.length || k <= 0)
        return new ArrayList<>();
    PriorityQueue<Integer> maxHeap = new PriorityQueue<>((o1, o2) -> o2 - o1);
    for (int num : nums) {
    
    
        maxHeap.add(num);
        if (maxHeap.size() > k)
            maxHeap.poll();
    }
    return new ArrayList<>(maxHeap);
}

quick selection

• Complexity: O(N) + O(1)
• Can only be used if modification of array elements is allowed

The partition() method of quick sort will return an integer j such that a[l...j-1] is less than or equal to a[j], and a[j+1...h] is greater than or equal to a[j]. At this time, a[j ] is the jth largest element of the array. This feature can be used to find the Kth element of the array. This algorithm for finding the Kth element is called the fast selection algorithm.

public ArrayList<Integer> GetLeastNumbers_Solution(int[] nums, int k) {
    
    
    ArrayList<Integer> ret = new ArrayList<>();
    if (k > nums.length || k <= 0)
        return ret;
    findKthSmallest(nums, k - 1);
    /* findKthSmallest 会改变数组，使得前 k 个数都是最小的 k 个数 */
    for (int i = 0; i < k; i++)
        ret.add(nums[i]);
    return ret;
}

public void findKthSmallest(int[] nums, int k) {
    
    
    int l = 0, h = nums.length - 1;
    while (l < h) {
    
    
        int j = partition(nums, l, h);
        if (j == k)
            break;
        if (j > k)
            h = j - 1;
        else
            l = j + 1;
    }
}

private int partition(int[] nums, int l, int h) {
    
    
    int p = nums[l];     /* 切分元素 */
    int i = l, j = h + 1;
    while (true) {
    
    
        while (i != h && nums[++i] < p) ;
        while (j != l && nums[--j] > p) ;
        if (i >= j)
            break;
        swap(nums, i, j);
    }
    swap(nums, l, j);
    return j;
}

private void swap(int[] nums, int i, int j) {
    
    
    int t = nums[i];
    nums[i] = nums[j];
    nums[j] = t;
}