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.
最基本的解法:sort默认升序排序返回nums[nums.size()-k],降序排序返回nums[k-1]。
class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
sort(nums.begin(), nums.end());
return nums[nums.size()-k];
}
};
显然这种做法是最常规最不讨面试官喜欢的做法。下面就要优化它。
(1)快排。
每次都可以将pivot元素放到它的最终位置上,这时就可以通过pivot的位置比k大还是比k小,来减少搜索的范围。
算法步骤:(假设第一大的元素在最左边)
1、初始化:left=0,right=nums.size()-1
2、Partition:如果pivot的位置正好是第k大即pos == k - 1,则返回pivot
3、如果pivot的位置在比k-1大,说明第k大的元素在它的左边,更新右边界值right = pos - 1
4、否则,说明第k大的元素在它的右边,更新左边界值eft = pos + 1
5、重复步骤2
class Solution {
public:
int partition(vector<int>& nums, int left, int right) {
int pivot = nums[left];
int l = left + 1, r = right;
while (l <= r) {
if (nums[l] < pivot && nums[r] > pivot)
swap(nums[l++], nums[r--]);
if (nums[l] >= pivot) l++;
if (nums[r] <= pivot) r--;
}
swap(nums[left], nums[r]);
return r;
}
int findKthLargest(vector<int>& nums, int k) {
int left = 0, right = nums.size() - 1;
while (true) {
int pos = partition(nums, left, right);
if (pos == k - 1) return nums[pos];
if (pos > k - 1) right = pos - 1;
else left = pos + 1;
}
}
};
(2)堆排。
为nums构建大顶堆。
class Solution {
private:
int heap_size;
public:
inline int left(int idx) {
return (idx << 1) + 1;
}
inline int right(int idx) {
return (idx << 1) + 2;
}
void max_heapify(vector<int>& nums, int idx) {
int largest = idx;
int l = left(idx), r = right(idx);
if (l < heap_size && nums[l] > nums[largest]) largest = l;
if (r < heap_size && nums[r] > nums[largest]) largest = r;
if (largest != idx) {
swap(nums[idx], nums[largest]);
max_heapify(nums, largest);
}
}
void build_max_heap(vector<int>& nums) {
heap_size = nums.size();
for (int i = (heap_size >> 1) - 1; i >= 0; i--)
max_heapify(nums, i);
}
int findKthLargest(vector<int>& nums, int k) {
build_max_heap(nums);
for (int i = 0; i < k; i++) {
swap(nums[0], nums[heap_size - 1]);
heap_size--;
max_heapify(nums, 0);
}
return nums[heap_size];
}
};
还可以使用优先队列:
class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
priority_queue<int> pq(nums.begin(), nums.end());
for (int i = 0; i < k - 1; i++)
pq.pop();
return pq.top();
}
};
或者用multiset:
class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
multiset<int> mset;
int n = nums.size();
for (int i = 0; i < n; i++) {
mset.insert(nums[i]);
if (mset.size() > k)
mset.erase(mset.begin());
}
return *mset.begin();
}
};