【Lintcode】544. Top k Largest Numbers

题目地址：

https://www.lintcode.com/problem/top-k-largest-numbers/description

求前K大元素返回，要求降序排列。直径用Quick Select即可。代码如下：

import java.util.Arrays;

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) {
        // write your code here
        if (nums == null || nums.length == 0) {
            return nums;
        }
        
        int pos = 0, l = 0, r = nums.length - 1;
        while ((pos = partition(nums, l, r)) != nums.length - k) {
            if (pos < nums.length - k) {
                l = pos + 1;
            } else {
                r = pos - 1;
            }
        }
        
        // 由于返回的时候要求降序排列，我们还得把前k大元素排列一下再返回。
        int[] res = new int[k];
        Arrays.sort(nums, nums.length - k, nums.length);
        for (int i = nums.length - k; i < nums.length; i++) {
            res[i - nums.length + k] = nums[2 * nums.length - k - 1 - i];
        }
        
        return res;
    }
    
    private int partition(int[] nums, int l, int r) {
        int m = l + (r - l >> 1);
        int tmp = nums[l];
        nums[l] = nums[m];
        nums[m] = tmp;
        
        int pivot = nums[l];
        while (l < r) {
            while (l < r && pivot <= nums[r]) {
                r--;
            }
            nums[l] = nums[r];
            while (l < r && pivot >= nums[l]) {
                l++;
            }
            nums[r] = nums[l];
        }
        
        nums[l] = pivot;
        return l;
    }
}

时间复杂度$O(n+k\log k)$。