LintCode 544. 前K大数

原题

解

用PriorityQueue的话，极度简单，以前的几道题已经做过无数遍了，直接上答案。

public class Solution {
    /*
     * @param nums: an integer array
     * @param k: An integer
     * @return: the top k largest numbers in array
     */
    public int[] topk(int[] nums, int k) {
        int[] result = new int[k];
        if (nums == null || nums.length < k) {
            return result;
        }
        
        Queue<Integer> pq = new PriorityQueue<>();
        for (int num : nums) {
            pq.add(num);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        
        for (int i = k - 1; i >= 0; i--) {
           result[i] = pq.poll(); 
        }
        
        return result;
    }
}

解2

用QuickSelect做。比普通QuickSelect要多一个步骤，因为这道题不是要第K大的数了，而是要前K个最大的数，同时还要排序。如果不排序，QuickSelect已经做到了，因为在找到第K大的数的时候，已经保证了比它大的都在一边了，那么那一边加上它本身，就是前K个了。题目要求是排序的，所以要再对这一部分进行一下排序。

public class Solution {
    /*
     * @param nums: an integer array
     * @param k: An integer
     * @return: the top k largest numbers in array
     */
    public int[] topk(int[] nums, int k) {
        int[] result = new int[k];
        if (nums == null || nums.length < k) {
            return result;
        }
        
        quickSort(nums, 0, nums.length - 1, k);
        // first k elements are the k largest elements
        quickSort(nums, 0, k - 1, -1);
        for (int i = 0; i < k; i++) {
            result[i] = nums[i];
        }
        
        return result;
    }
    
    private void quickSort(int[] nums, int start, int end, int k) {
        if (start >= end) {
            return;
        }
        
        int left = start;
        int right = end;
        int pivot = nums[start + (end - start) / 2];
        while (left <= right) {
            while (left <= right && nums[left] > pivot) {
                left++;
            }
            
            while (left <= right && nums[right] < pivot) {
                right--;
            }
            
            if (left <= right) {
                int temp = nums[left];
                nums[left] = nums[right];
                nums[right] = temp;
                
                left++;
                right--;
            }
        }
        
        if (k != -1) {
            if (start + k - 1 < right + 1) {
                quickSort(nums, start, right, k);
            } else if (start + k - 1 > right + 1) {
                quickSort(nums, left, end, k - (left - start));
            }
        } else {
            quickSort(nums, start, right, k);
            quickSort(nums, left, end, k);
        }
    }
}

这里我为了省事，直接用QuickSort再sort了一遍，因为QuickSelect跟QuickSort的第一步都是一样的，我用了一个-1当作k的特殊值，如果是-1就用QuickSort，不是就QuickSelect。

分析

第一个方法，时间是O(nlogk)，空间是O(k)。第二个方法，基于QuickSelect，空间常数O(1)，时间第一步是O(n)，再QuickSort前k个数字，即O(klogk)，所以总共就是O(n+klogk)。这也很合理，如果k等于n，那么就等于sort了一遍array。