[leetcode] 703. Kth Largest Element in a Stream @ python

最新推荐文章于 2021-08-12 21:54:08 发布

原创最新推荐文章于 2021-08-12 21:54:08 发布 · 650 阅读

CC 4.0 BY-SA版权

本文介绍了一种在数据流中实时查找第K大元素的算法，提出了两种解决方案：一是利用二分查找法，对初始数据进行排序，然后在每次插入新元素时维护排序状态；二是使用小根堆，通过保持堆大小等于K来确保堆顶元素即为第K大元素。这两种方法在处理大量数据流时表现出高效的时间复杂度。

原题

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Your KthLargest class will have a constructor which accepts an integer k and an integer array nums, which contains initial elements from the stream. For each call to the method KthLargest.add, return the element representing the kth largest element in the stream.

Example:

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
Note:
You may assume that nums’ length ≥ k-1 and k ≥ 1.

解法1

二分查找法. 初始化时将nums升序排列, 然后每次往里加数字时, 使用bisect.insort_right将val加到指定位置, 如果指定位置已有相同值, 则往右边加. 最后返回index为len(self.data) - self.k的值即可.
Time: O(log(n))
Space: O(1)

代码

class KthLargest:

    def __init__(self, k: 'int', nums: 'List[int]'):
        self.data = sorted(nums)
        self.k = k
        

    def add(self, val: 'int') -> 'int':
        bisect.insort_right(self.data, val)
        return self.data[len(self.data) - self.k]
        


# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)

解法2

小根堆. https://docs.python.org/3.0/library/heapq.html
由于我们需要返回第K大的值, 因此当nums的长度大于K时, 将多余的元素删除, 这样当nums长度为K时, 求第K大的值就是heapq里的最小值.

这里要注意edge case, 当初始化nums的长度小于K时, 由于self.data已转化为heapq, 第一个值就是最小值, 此时我们将val加入到self.data, 然后返回self.data[0]即可.

当nums的长度已经等于K且val比最小值大时, 使用heapq.heapreplace将val替代最小值, 然后返回self.data[0].

代码

class KthLargest(object):

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        self.data = nums
        self.k = k
        heapq.heapify(self.data)
        while len(self.data) > self.k:
            heapq.heappop(self.data)
        

    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        if len(self.data) < self.k:
            heapq.heappush(self.data, val)
        elif val > self.data[0]:
            heapq.heapreplace(self.data, val)
        return self.data[0]


# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)
        


# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)