一种解决方案是使用「三路划分」,即每轮将数组划分为三个部分:小于、等于和大于基准数的所有元素。这样当发现第 k 大数字处在"等于基准数"的子数组中时,便可以直接返回该元素。
为了进一步提升算法的稳健性,我们采用随机选择的方式来选定基准数。
代码实现(快排):
复制代码
public class Solution {
private int quickSelect(List<Integer> nums, int k) {
// 随机选择基准数
Random rand = new Random();
int pivot = nums.get(rand.nextInt(nums.size()));
// 将大于、小于、等于 pivot 的元素划分至 big, small, equal 中
List<Integer> big = new ArrayList<>();
List<Integer> equal = new ArrayList<>();
List<Integer> small = new ArrayList<>();
for (int num : nums) {
if (num > pivot)
big.add(num);
else if (num < pivot)
small.add(num);
else
equal.add(num);
}
// 第 k 大元素在 big 中,递归划分
if (k <= big.size())
return quickSelect(big, k);
// 第 k 大元素在 small 中,递归划分
if (nums.size() - small.size() < k)
return quickSelect(small, k - nums.size() + small.size());
// 第 k 大元素在 equal 中,直接返回 pivot
return pivot;
}
public int findKthLargest(int[] nums, int k) {
List<Integer> numList = new ArrayList<>();
for (int num : nums) {
numList.add(num);
}
return quickSelect(numList, k);
}
}
代码实现(堆排):
复制代码
class Solution {
public:
void adjMinHeap(vector<int>& nums, int root, int heapsize) {
int left = root * 2 + 1, right = root * 2 + 2, minimum = root;
if (left < heapsize && nums[left] < nums[minimum])
minimum = left;
if (right < heapsize && nums[right] < nums[minimum])
minimum = right;
if (minimum != root) {
swap(nums[minimum], nums[root]);
adjMinHeap(nums, minimum, heapsize);
}
}
void buildMinHeap(vector<int>& nums, int k) {
for (int i = k / 2 - 1; i >= 0; i--)
adjMinHeap(nums, i, k);
}
int findKthLargest(vector<int>& nums, int k) {
buildMinHeap(nums, k);
for (int i = k; i < nums.size(); i++) {
if (nums[i] < nums[0])
continue;
swap(nums[0], nums[i]);
adjMinHeap(nums, 0, k);
}
return nums[0];
}
};