python 无序列表中第k大元素_LeetCode-python 703.数据流中的第k大元素

题目链接

难度：简单       类型： 堆

设计一个找到数据流中第K大元素的类(class)。注意是排序后的第K大元素，不是第K个不同的元素。

你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器，它包含数据流中的初始元素。每次调用 KthLargest.add，返回当前数据流中第K大的元素。

示例

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

kthLargest.add(10); // returns 5

kthLargest.add(9); // returns 8

kthLargest.add(4); // returns 8

解题思路

维持一个k个元素的最小堆

当新元素小于堆顶元素时，不入堆

当新元素大于堆顶元素时，替换堆顶元素，再变换成最小堆

代码实现

import heapq

class KthLargest:

def __init__(self, k, nums):

"""

:type k: int

:type nums: List[int]

"""

self.heap = nums

self.k = k

heapq.heapify(self.heap)

while len(self.heap) > k:

heapq.heappop(self.heap)

def add(self, val):

"""

:type val: int

:rtype: int

"""

if len(self.heap) < self.k:

heapq.heappush(self.heap, val)

elif val > self.heap[0]:

heapq.heapreplace(self.heap, val)

return self.heap[0]