215. Kth Largest Element in an Array [Medium]

经典的快速选择算法题目

快速选择算法的本质是快速排序算法，但是每确定一个元素的位置，下一次就只需要在其中一侧寻找，而且找到第k大元素就可以停止，比完整排序的效率高一些

快速排序算法最优、平均时间复杂度为O(nlogn) ，最差时间复杂度为O(n^2)

快速选择算法最优、平均时间复杂度为O(n)，最差时间复杂度为O(n^2)

以下代码全都是快速选择，很多种不同的写法

/**
 * 快速选择算法，平均时间复杂度O(n)，空间复杂度O(logn)
 * Runtime: 7 ms, faster than 25.36%
 * Memory Usage: 39.3 MB, less than 54.19%
 */
class Solution {
    public int findKthLargest(int[] nums, int k) {
        return helper(nums, 0, nums.length - 1, k);
    }
    private int helper(int[] nums, int start, int end, int k) { // 用首尾指针且保存轴值的方法找到pivot的位置（最后没有将pivot值赋给找到的位置是因为在本题中不需要，pivot两侧的数据已经没错了，但排序会需要）
        int left = start, right = end, pivot = nums[right];
        while (left < right) {
            while (left < right && nums[left] >= pivot)
                left++;
            nums[right] = nums[left];
            while (left < right && nums[right] <= pivot)
                right--;
            nums[left] = nums[right];
        }
        if (left == k - 1) // 找到第k大
            return pivot;
        if (left > k - 1) // 第k大在pivot的左侧
            return helper(nums, start, --left, k);
        return helper(nums, ++right, end, k); // 第k大在pivot的右侧
    }
}

/**
 * 快速选择算法
 * Runtime: 16 ms, faster than 13.92%
 * Memory Usage: 39.7 MB, less than 11.99%
 */
class Solution {
    public int findKthLargest(int[] nums, int k) {
        return helper(nums, 0, nums.length - 1, k);
    }
    private int helper(int[] nums, int start, int end, int k) { // 用两个都从头部出发的指针找到pivot的位置
        int pivotIdx = start, pivot = nums[end];
        for (int i = start; i < end; i++) {
            if (nums[i] >= pivot)
                swap(nums, i, pivotIdx++);
        }
        swap(nums, pivotIdx, end);
        if (pivotIdx == k - 1)
            return pivot;
        if (pivotIdx > k - 1)
            return helper(nums, start, pivotIdx - 1, k);
        return helper(nums, pivotIdx + 1, end, k);
    }
    private void swap(int[] nums, int a, int b) {
        int temp = nums[a];
        nums[a] = nums[b];
        nums[b] = temp;
    }
}

/**
 * 上一段代码的非递归版本
 * Runtime: 19 ms, faster than 7.94%
 * Memory Usage: 39.5 MB, less than 30.96%
 */
class Solution {
    public int findKthLargest(int[] nums, int k) {
        int start = 0, end = nums.length - 1;
        while (start <= end) { // 用这个循环替代递归
            int pivotIdx = start, pivot = nums[end];
            for (int i = start; i < end; i++) {
                if (nums[i] >= pivot)
                    swap(nums, i, pivotIdx++);
            }
            swap(nums, pivotIdx, end);
            if (pivotIdx == k - 1)
                return pivot;
            if (pivotIdx > k - 1) {
                end = pivotIdx - 1;
            } else {
                start = pivotIdx + 1;
            }
        }
        return -1;
    }
    private void swap(int[] nums, int a, int b) {
        int temp = nums[a];
        nums[a] = nums[b];
        nums[b] = temp;
    }
}

/**
 * 非递归，另一种写法
 * Runtime: 8 ms, faster than 22.62%
 * Memory Usage: 39.3 MB, less than 41.43%
 */
class Solution {
    public int findKthLargest(int[] nums, int k) {
        int start = 0, end = nums.length - 1;
        while (start <= end) {
            int idx = findPivotIdx(nums, start, end);
            if (idx == k - 1)
                return nums[idx];
            if (idx > k - 1) {
                end = idx - 1;
            } else {
                start = idx + 1;
            }
        }
        return -1;
    }
    private int findPivotIdx(int[] nums, int start, int end) {
        int pivotIdx = start, pivot = nums[start];
        while (start <= end) {
            while (start <= end && nums[start] >= pivot) {
                start++;
            }
            while (start <= end && nums[end] <= pivot) {
                end--;
            }
            if (start > end)
                break;
            swap(nums, start, end);
        }
        swap(nums, pivotIdx, end);
        return end;
    }
    private void swap(int[] nums, int a, int b) {
        int temp = nums[a];
        nums[a] = nums[b];
        nums[b] = temp;
    }
}

第二次做：

也是用快速选择算法，但因为每找到一个轴值位置，可以将接下来查找的范围限定在它的一侧

所以可以不用辅助函数，直接用一个while循环，每次更新查找范围（start和end），这样的效率高了很多

/**
 * Runtime: 1 ms, faster than 97.73%
 * Memory Usage: 39 MB, less than 89.13%
 */
class Solution {
    public int findKthLargest(int[] nums, int k) {
        int start = 0, end = nums.length - 1, pivotIdx;
        while (start < end) {
            int midIdx = (start + end) / 2;
            int pivot = nums[midIdx], left = start, right = end;
            nums[midIdx] = nums[right]; // 将最右侧的值保存到选取的轴值位置
            while (left < right) { // 结束后left = right等于轴值的index
                while (left < right && nums[left] <= pivot) {
                    left++;
                }
                nums[right] = nums[left];
                while (left < right && nums[right] > pivot) {
                    right--;
                }
                nums[left] = nums[right];
            }
            nums[left] = pivot; // 将轴值更新到正确的位置
            if (left == nums.length - k) {
                return nums[left];
            }
            if (left > nums.length - k) {
                end = right - 1;
            } else {
                start = left + 1;
            }
        }
        return nums[start]; // 因为由题意，第k大的数一定存在，否则应该return start == nums.length - k ? nums[start] : -1;
    }
}