Kth Largest Element in an Array

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

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

For example,

Given [3,2,1,5,6,4] and k = 2, return 5.

Note:

You may assume k is always valid, 1 ≤ k ≤ array's length.

Thoughts

1. 是找第k大的，因此是倒序。

2. Partition: 把pivot放到应置的位置并返回对应的位置。用双头指针通过swap把比它大的放前面，剩下的放后面。

3. 通过每次返回的排好的位置pos和K作比较，此时pos前或pos后的都比K小或大从而（类似二分法）调整双头指针（start和end）位置抛去pos前或后部分再做partition. 直到pos刚好==k - 1。

Code

/*
 * @lc app=leetcode id=215 lang=cpp
 *
 * [215] Kth Largest Element in an Array
 */

// @lc code=start
class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        const auto partition = [&](int l, int r) {
            swap(nums[l], nums[rand() % (r - l + 1) + l]);
            const auto s = l, t = nums[l++];
            while (l <= r) {
                if (nums[l] > t) ++l;
                else if (nums[r] <= t) --r;
                else {
                    swap(nums[l++], nums[r--]);
                } 
            }
            swap(nums[s], nums[r]);
            return r;
        };
        int start = 0, end = nums.size() - 1;
        --k;
        while (true) {
            int pos = partition(start, end);
            if (pos == k) return nums[pos];
            else if (pos < k) start = pos + 1;
            else end = pos - 1;
        } 
        return nums[0];
    }
};
// @lc code=end

/*
 * @lc app=leetcode id=215 lang=cpp
 *
 * [215] Kth Largest Element in an Array
 *
 * https://leetcode.com/problems/kth-largest-element-in-an-array/description/
 *
 * algorithms
 * Medium (49.42%)
 * Likes:    2338
 * Dislikes: 189
 * Total Accepted:    424.2K
 * Total Submissions: 857.4K
 * Testcase Example:  '[3,2,1,5,6,4]\n2'
 *
 * Find the kth largest element in an unsorted array. Note that it is the kth
 * largest element in the sorted order, not the kth distinct element.
 * 
 * Example 1:
 * 
 * 
 * Input: [3,2,1,5,6,4] and k = 2
 * Output: 5
 * 
 * 
 * Example 2:
 * 
 * 
 * Input: [3,2,3,1,2,4,5,5,6] and k = 4
 * Output: 4
 * 
 * Note: 
 * You may assume k is always valid, 1 ≤ k ≤ array's length.
 * 
 */
class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        int start = 0, end = nums.size() - 1;
        const auto partition = [&]() {
            int pivot = nums[start], l = start + 1, r = end;
            // 确保r在到了l处还会接着往前走跨过所有和pivot相等的元素
            while (l <= r) {
                // l 停在第一个比p小的位置
                if (nums[l] >= pivot) ++l;
                // r 停在第一个比p大的位置
                else if (nums[r] <= pivot) --r;
                else swap(nums[l], nums[r]);
            }
            // 倒序， start 位置要放比p大的
            swap(nums[start], nums[r]);
            return r;
        };
        while (true) {
            const auto pos = partition();
            if (pos == k - 1) return nums[pos];
            else if (pos < k - 1) start = pos + 1;
            else end = pos - 1;
        }
    }
};

Analysis

平均来说, 每次partition把原数组分成了前后各一半(pos == mid), 下次搜索时问题规模减半, 因此T(n) = T(n/2) + O(n) , O(n) 用于partition, 时间复杂度O(N). 最差时每次partition只能减1个问题规模(逆序时), 时间复杂度O(N^2).

Ver.2

还可以用heap sort做, 维持大小为k的min heap, 此时heap中顶元素即第k大.

class Solution {
    public int findKthLargest(int[] nums, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<>();
        for (int num : nums) {
            pq.offer(num);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        int res = 0;
        if (!pq.isEmpty()) {
            res = pq.peek();
        }
        return res;
    }
}