王廷瑋|數位醫療|智慧醫療: 215. Kth Largest Element in an Array WFU

2024年7月18日 星期四

215. Kth Largest Element in an Array

215. Kth Largest Element in an Array


給定一個整數數組 nums 和一個整數 k,返回數組中第 k 大的元素。

注意,這裡指的是排序後的第 k 大的元素,而不是第 k 個不同的元素。

你能不使用排序來解決這個問題嗎?

範例:

輸入:nums = [3,2,1,5,6,4], k = 2 輸出:5


Python


import heapq

class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
# Create a min heap
heap = []
# Process each number in the array
for num in nums:
# If the heap size is less than k, just add the number
if len(heap) < k:
heapq.heappush(heap, num)
# If the current number is greater than the smallest in the heap,
# remove the smallest and add the current number
elif num > heap[0]:
heapq.heapreplace(heap, num)
# The root of the heap is the kth largest element
return heap[0]

29.52MB, 467ms


C++


class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
// Create a min heap
priority_queue<int, vector<int>, greater<int>> pq;
// Process each number in the array
for (int num : nums) {
// If the heap size is less than k, just add the number
if (pq.size() < k) {
pq.push(num);
}
// If the current number is greater than the smallest in the heap,
// remove the smallest and add the current number
else if (num > pq.top()) {
pq.pop();
pq.push(num);
}
}
// The top of the heap is the kth largest element
return pq.top();
}
};

61.54MB, 87ms


Javascript


/**
* @param {number[]} nums
* @param {number} k
* @return {number}
*/
var findKthLargest = function(nums, k) {
// Create a min heap
let heap = new MinHeap();
// Process each number in the array
for (let num of nums) {
// If the heap size is less than k, just add the number
if (heap.size() < k) {
heap.push(num);
}
// If the current number is greater than the smallest in the heap,
// remove the smallest and add the current number
else if (num > heap.peek()) {
heap.pop();
heap.push(num);
}
}
// The root of the heap is the kth largest element
return heap.peek();
};

// MinHeap implementation
class MinHeap {
constructor() {
this.heap = [];
}
push(val) {
this.heap.push(val);
this.bubbleUp(this.heap.length - 1);
}
pop() {
const min = this.heap[0];
const last = this.heap.pop();
if (this.heap.length > 0) {
this.heap[0] = last;
this.bubbleDown(0);
}
return min;
}
peek() {
return this.heap[0];
}
size() {
return this.heap.length;
}
bubbleUp(index) {
while (index > 0) {
const parentIndex = Math.floor((index - 1) / 2);
if (this.heap[parentIndex] <= this.heap[index]) break;
[this.heap[parentIndex], this.heap[index]] = [this.heap[index], this.heap[parentIndex]];
index = parentIndex;
}
}
bubbleDown(index) {
while (true) {
let minIndex = index;
const leftChild = 2 * index + 1;
const rightChild = 2 * index + 2;
if (leftChild < this.heap.length && this.heap[leftChild] < this.heap[minIndex]) {
minIndex = leftChild;
}
if (rightChild < this.heap.length && this.heap[rightChild] < this.heap[minIndex]) {
minIndex = rightChild;
}
if (minIndex === index) break;
[this.heap[index], this.heap[minIndex]] = [this.heap[minIndex], this.heap[index]];
index = minIndex;
}
}
}

67.05MB, 101ms