Kth Largest Element in an Array

Description

Given an integer array nums and an integer k, return the kth largest element in the array.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

You must solve it in O(n) time complexity.

Example 1:

Input: nums = [3,2,1,5,6,4], k = 2 Output: 5

Example 2:

Input: nums = [3,2,3,1,2,4,5,5,6], k = 4 Output: 4

Constraints:

• 1 <= k <= nums.length <= 105
• -104 <= nums[i] <= 104

Code

Quick Select

Time Complexity: $O(n)$, Space Complexity: $O(n)$

• Partial sorting wiki
• std::random_device
• std::mt19937
• std::shuffle

若沒有 randonmized input，runtime 會慢很多（240ms），有的話只需要（102ms）。

• Sorting
• Quick Sort

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        // randonmize input to gaurantee O(n)
        random_device rd;
        mt19937 g(rd());
        shuffle(nums.begin(), nums.end(), g);
        int left = 0, right = nums.size() - 1;
        k = nums.size() - k;
        while(left < right) {
            int index = quick_select(nums, left, right);
            if(index == k) return nums[index];
            else if(index < k) left = index + 1;
            else right = index - 1;
        }
 
        return nums[left];
 
 
    }
 
    int quick_select(vector<int>& nums, int front, int end) {
        int pivot = nums[end];
        int i = front - 1;
        for(int j = front; j < end; j++) {
            if(nums[j] < pivot) {
                i++;
                swap(nums[j], nums[i]);
            }
        }
        i++;
        swap(nums[i], nums[end]);
        return i;
    }
};

partial sorting(nth_element & partial_sort)

Time Complexity: $O(n)$, Space Complexity: $O(n)$

• std::nth_element

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        nth_element(nums.begin(), nums.begin() + k - 1, nums.end(), greater<int>());
        return nums[k - 1];
    }
};

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        nth_element(nums.begin(), nums.begin() + k - 1, nums.end(), [](const int a, const int b){return a > b;});
        return nums[k - 1];
    }
};

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        k = nums.size() - k;
        nth_element(nums.begin(), nums.begin() + k, nums.end());
        return nums[k];
    }
};

• std::partial_sort

class Solution {
public:
 
    int findKthLargest(vector<int>& nums, int k) {
        partial_sort(nums.begin(), nums.begin() + k, nums.end(), greater<int>());
        return nums[k - 1];
    }
};

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        k = nums.size() - k;
        partial_sort(nums.begin(), nums.begin() + k + 1, nums.end());
        return nums[k];
    }
};

priority_queue default 為 max heap，而 multiset default 為 min heap。

在寫 custom cmp 時，記得只有 priority_queue 的大小方向和其他人相反。

Max Heap

Time Complexity: $O(n\log k)$, Space Complexity: $O(k)$

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int> p;
        for(auto n: nums) {
            p.push(n);
        }
 
        for(int i = 0; i < k - 1; i++) {
            p.pop();
        }
 
        return p.top();
    }
};

注意 multiset 的 cmp 大小比較順序。

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        auto cmp = [](const int a, const int b) {return a > b;};
        multiset<int, decltype(cmp)> m(cmp);
        for(auto n: nums) {
            m.insert(n);
        }
 
        for(int i = 0; i < k - 1; i++) m.erase(m.begin());
        return *m.begin();
    }
};

Min Heap

Time Complexity: $O(n\log k)$, Space Complexity: $O(k)$

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int>> min_heap;
        for(auto& n: nums) {
            min_heap.push(n);
            if(min_heap.size() > k) 
                min_heap.pop();
        }
 
        return min_heap.top();
    }
 
};

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        auto cmp = [](const int a, const int b){return a > b;};
        priority_queue<int, vector<int>, decltype(cmp)> p(cmp);
        for(auto n: nums) {
            p.push(n);
            if(p.size() > k)
                p.pop();
        }
 
        return p.top();
    }
};

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        multiset<int> m;
        for(auto n: nums) {
            m.insert(n);
            if(m.size() > k) {
                m.erase(m.begin());
            }
        }
 
        return *m.begin();
    }
};

Counting Sort

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        
        int size = 20005;
        vector<int> count(size, 0);
 
        for(auto n: nums) {
            int index = n + 10000;
            count[index]++;
        }
 
        for(int i = 1; i < count.size(); i++) {
            count[i] += count[i - 1];
        }
 
        int n = nums.size();
        vector<int> sorted(n);
        for(int i = n - 1; i >= 0; i--) {
            sorted[count[nums[i] + 10000] - 1] = nums[i];
            count[nums[i] + 10000]--;
        }
 
        k = n - k;
        return sorted[k];
    }
 
 
};

Source

• Kth Largest Element in an Array - LeetCode