[LintCode][Sort][Heap] Kth Largest Element

Problem

More Discussions
Find K-th largest element in an array.

Example

In array [9,3,2,4,8], the 3rd largest element is 4.

In array [1,2,3,4,5], the 1st largest element is 5, 2nd largest element is 4, 3rd largest element is 3 and etc.

Challenge

O(n) time, O(1) extra memory.

Solution

方法一：用最小堆。存k个数，如果当前的数大于最小堆堆顶的数，则弹出堆顶数，当前数入堆。时间O(nlogk)

class Solution {
public:
    /*
     * param k : description of k
     * param nums : description of array and index 0 ~ n-1
     * return: description of return
     */
    int kthLargestElement(int k, vector nums) {
        // write your code here
        priority_queue, greater > minHeap;
        for(int i = 0; i < nums.size(); i++) {
            if (minHeap.size() < k) {
                minHeap.push(nums[i]);
            } else {
                if (nums[i] > minHeap.top()) {
                    minHeap.pop();
                    minHeap.push(nums[i]);
                }
            }
        }
        
        return minHeap.top();
    }
};

方法二：快排的思想，找到第k个数。空间O(1)，时间的话平均复杂度O(n)。

class Solution {
public:
    /*
     * param k : description of k
     * param nums : description of array and index 0 ~ n-1
     * return: description of return
     */
    void partition(vector &a, int beg, int end, int k) {
        if (beg == end) {
            return;
        }
        
        int index = rand() % (end - beg) + beg;
        swap(a[index], a[beg]);
        int key = a[beg];
        int i = beg;
        int j = end;
        while (i < j) {
            if (a[j] >= key) {
                j--;
            } else {
                i++;
                swap(a[i], a[j]);
            }
        }
        
        int mid = i;
        swap(a[beg], a[i]);
        
        if (mid == k) {
            return;
        } else if (mid > k) {
            partition(a, beg, mid - 1, k);
        } else {
            partition(a, mid + 1, end, k);
        }
    }
    
    int kthLargestElement(int k, vector nums) {
        srand (time(NULL));
        partition(nums, 0, nums.size() - 1, nums.size() - k);
        
        return nums[nums.size() - k];
    }
};