算法笔记 -- 快速选择算法

这个算法的来源：215. 数组中的第K个最大元素。

今天在写leetcode的时候，写到了这题，也是发现了一个快排的新用途，可以使用快排的分治过程来寻找topK。

最初的思路是，对数组排序，然后取下标为n-k即第k大的元素，算法的思路也比较简单。

JAVA 复制

class Solution {
    public int findKthLargest(int[] nums, int k) {
        Arrays.sort(nums);
        return nums[nums.length - k];
    }
}

时间复杂度为$O(nlogn)$，空间复杂度为$O(1)$。

另外，除了排序，还可以使用堆来优化，因为我们只需要求第K个，因此可以创造一个大顶堆，先将所有元素都添加进去，然后依次出队k个即可。

JAVA 复制

class Solution {
    public int findKthLargest(int[] nums, int k) {
        Queue<Integer> q = new PriorityQueue<>((n1, n2) -> {return n2 - n1;});
        for(int num : nums) {
            q.offer(num);
        }
        while(q.size() != nums.length - k + 1) {
            q.poll();
        }
        return q.poll();
    }
}

后面发现，自己的这个解法还是不够优雅，其实不必将所有的元素都添加进堆中，只需要在堆中保存k个元素。堆的话，我们选择小顶堆，因为这样取第k个元素就只需要取队列的peek即可。

遍历一趟nums时，若堆中元素个数小于k个时，直接入堆即可，当堆的元素个数大于k时，判读此时的元素和堆顶的大小，如果当前堆顶元素比当前的入堆元素大，那么说明这个数显然不会是第k个数及它后面的数，跳过即可，如果当前的元素比堆顶元素大，将堆顶出队，将当前元素入队，直到遍历完所有的，此时堆顶即时我们所求的第k大的元素。

JAVA 复制

class Solution {
    public int findKthLargest(int[] nums, int k) {
        Queue<Integer> q = new PriorityQueue<>();
        for(int num : nums){
            if(q.size() < k) {
                q.offer(num);
            } else {
                if(q.peek() > num) {
                    continue;
                } else {
                    q.poll();
                    q.offer(num);
                }
            }
        }
        return q.peek();
    }
}

时间复杂度为$O(nlogk)$，空间复杂度为$O(k)$。

接下里就是了解到的第三种算法，也是听过很多次的快速选择算法，可以做到在一个随机序列中，达到期望值为$O(n)$的时间复杂度。

快速选择算法的思路主要如下：

1. 随机选择一个基准值x，遍历一趟序列，确保完成这个过程后，在x左边的元素都$\le x$，而在x右边的元素都$\ge x$；
2. 完成上述过程后，可以确保x已经在其排序的位置上，如果x所处位置为所求位置k，那么返回即可，否则，比较两边的元素数组，左边元素数量为当前x所处位置 - 起始点，若$leftLength<k$，则去右边找，否则继续在左边找，避免了快速排序中两次递归的过程；

JAVA 复制

class Solution {
    public int findKthLargest(int[] nums, int k) {
        return quickSelect(nums, 0, nums.length - 1, nums.length - k + 1);
    }
    private int quickSelect(int[] nums, int l, int r, int k) {
        if(l >= r) {
            return nums[l];
        }
        int x = nums[l + r >> 1], i = l - 1, j = r + 1;
        while(i < j) {
            do i++; while(nums[i] < x);
            do j--; while(nums[j] > x);
            if(i < j) {
                int tmp = nums[i];
                nums[i] = nums[j];
                nums[j] = tmp;
            }
        }
        // 左边长度比k小,在左边寻找
        if(k <= j - l + 1) return quickSelect(nums, l, j, k);
        // 否则去右边寻找
        else return quickSelect(nums, j + 1, r, k - (j - l + 1));
    }
}

时间复杂度为$O(n)$，空间复杂度为$O(logn)$。

同样的，由于我们可以发现，其实递归的次数会比较小，在$logn$级别，因此可以直接使用while循环替换即可。

JAVA 复制

private int quickSelect(int[] nums, int l, int r, int k) {
    while(true) {
        if(l == r) {
            return nums[l];
        }
        int i = l - 1, j = r + 1;
        while(i < j) {
            do i++; while (nums[i] > x);
            do j--; while (nums[j] < x);
            if(i < j) {
                int t = nums[i];
                nums[i] = nums[j];
                nums[j] = t;
            }
        }
        if(k <= j - l + 1) {
            r = j;
        } else {
            k -= (j - l + 1);
            l = j + 1;
        }
    }
}

时间复杂度为$O(n)$，空间复杂度为$O(1)$。