题目描述

在未排序的数组中找到第 k 个最大的元素。请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

示例 1:

输入: [3,2,1,5,6,4] 和 k = 2
输出: 5

示例 2:

输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
输出: 4

说明:

你可以假设 k 总是有效的,且 1 ≤ k ≤ 数组的长度。

https://leetcode-cn.com/problems/kth-largest-element-in-an-array/

解法1 – 排序O(nlogn)

这是一道经典的题目,有多种解法。最直接的想法,我们可以先排序,然后取出第n – k个数字。时间复杂度O(nlogn),空间复杂度O(1)。全部代码如下:

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

运行时间:7ms

解法2 – 堆(优先队列)O(nlogk)

我们维护一个大小为k的最小堆,这可以通过Java的PriorityQueue的默认构造方法构建。我们不断的往堆中加入元素,直到元素数量超过k,我们删除最小元素。当扫描完整个nums数组,堆中就留下了k个最大的数,我们取k个最大数中最小的那一个就是答案。

例如nums = {3,2,1,5,6,4}, k = 2。当堆内容为{3, 2}时,再加入1得到{3, 2, 1},元素数量超过了k,我们删除最小元素1得到{3, 2};再加入5得到{3, 2, 5},删除最小元素2得到{3, 5};再加入6得到{3, 5, 6},删除最下元素3得到{5, 6};再加入4得到{5, 6, 4},删除最下元素4得到{5, 6}。此时,我们从堆{5, 6}取出最小元素5就是答案。

根据Java doc,对于offerpoll操作的时间复杂度是O(logn)。堆的大小为k,我们遍历数组每个元素,时间复杂度为O(nlogk),空间复杂度为O(k)。

class Solution {
    public int findKthLargest(int[] nums, int k) {
        PriorityQueue<Integer> queue = new PriorityQueue<>(k);

        for (int num : nums) {
            queue.offer(num);
            if (queue.size() > k)
                queue.poll();
        }

        return queue.peek();
    }
}

运行时间:14ms。理论上解法2的时间复杂度比解法1低,但是运行时间却比解法1长。我猜测是因为需要开辟额外空间,box/unbox造成的开销导致的。如果有不对的请指正。

解法3 – Partition函数 O(n)

解法3使用快速排序的Partition函数。Partition函数执行完会返回pivot所在的索引i,满足i左边的元素都小于等于pivot,i右边的元素都大于pivot。为了方便思考与实现,我们搜索第k小的数(k从1开始)。如果需要搜索第k大的数字,我们可以求第|nums| – k + 1小的数字,就是第k大的数字。

我们根据pivot所在的索引i与k的关系递归的调用Partition函数,我们将Partition函数命名为findKth,函数签名为int findKth(int[] nums, int start, int end, int k)。如果i == k – 1,那么我们就找到了第k大的数(因为k从1起,我们需要-1);如果i > k – 1,我们将搜索范围缩小到[start, i – 1]继续搜索;如果i < k – 1,我们将搜索范围缩小到[i+1, end]继续搜索;如果start>=end表示已经越界,我们返回nums[start]。

该算法最优的时间复杂度为O(n),空间复杂度为O(1)。

证明:第一次访问n个元素,第二次访问n/2个元素;第三次访问n/4个元素……。
循环次数C = n + n/2 + n/4 + n/8 + n/16 + ... = n * (1 + 1/2 + 1/4 + 1/8 + 1/16 + ...)
其中1/2 + 1/4 + 1/8 + 1/16 + ...是无穷级数,收敛于1。那么C = n * (1 + 1) = 2*n。
因此时间复杂度为O(n)。

但是需要注意的是这是最优的情况,如果nums数组完全有序,那么第一次访问n个元素,第二次访问n-1个元素,第三次访问n-2个元素……。很明显,这种case的时间复杂度为O(n^2)。为了避免这种情况,我们可以把nums数组随机打乱。随机打乱算法的时间复杂度为O(n)。

从直觉上来讲,我们将数组nums随机打乱需要扫描一边数组,这会不会使得运行时间增加?是的,如果nums数组分布比较随机,那么随机打乱是会增加运行时间的。但是如果nums数组有序,在O(n^2)的时间复杂度下运行时间远远超出随机打乱的开销。

这是两次运行时间的对比:
    1. 不随机打乱运行时间30ms
    2. 随机打乱运行时间6ms
图1 – 运行时间对比

全部代码如下:

class Solution {
    public int findKthLargest(int[] nums, int k) {
        shuffle(nums);
        return findKth(nums, 0, nums.length - 1, nums.length - k + 1);
    }

    private void shuffle(int[] nums) {
        Random random = new Random();
        for (int i = 1; i < nums.length; i++) {
            int j = random.nextInt(i + 1);
            int tmp = nums[i];
            nums[i] = nums[j];
            nums[j] = tmp;
        }
    }

    // k start with 1
    int findKth(int[] nums, int start, int end, int k) {
        if (start >= end)
            return nums[start];
        int pivot = nums[end];
        int i = start;
        for (int j = start; j < end; j++) {
            if (nums[j] <= pivot) {
                int tmp = nums[i];
                nums[i] = nums[j];
                nums[j] = tmp;
                i++;
            }
        }

        int tmp = nums[i];
        nums[i] = nums[end];
        nums[end] = tmp;

        if (i == k - 1)
            return nums[i];
        else if (i > k - 1)
            return findKth(nums, start, i - 1, k);
        else
            return findKth(nums, i + 1, end, k);
    }
}

参考:

  1. https://leetcode.com/problems/kth-largest-element-in-an-array/discuss/60294/Solution-explained
  2. https://rcoh.me/posts/linear-time-median-finding/
  3. https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle

Leave a Reply

Your email address will not be published. Required fields are marked *