leetcode703.数据流中的第K大元素

88 阅读1分钟

703.数据流中的第K大元素

设计一个找到数据流中第K大元素的类(class)。注意是排序后的第K大元素,不是第K个不同的元素。你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器,它包含数据流中的初始元素。每次调用 KthLargest.add,返回当前数据流中第K大的元素。 nums 的长度≥ k-1k ≥ 1。

int k = 3;	int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.heap = nums
        heapq.heapify(self.heap)  # 原地转化为小顶堆

    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)
        while len(self.heap) > self.k:  # 长度超过K,维护堆
            heapq.heappop(self.heap)  # 弹出并返回 heap 的最小的元素
        return self.heap[0]  # 返回 heap 的最小的元素
# 想找到第K大元素,则维护K个元素的MinHeap;
# 想找到第K小元素,则维护K个元素的MaxHeap。
# 堆顶元素(最小元素)是第K大元素。