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

原题

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)