日野弥生：勉強しよう
LeetCode 703 - 数据流中的第 K 大元素
发表于2025年03月20日
一个很显而易见的解法是，先将数组降序排列好，然后返回数组中第k个数。
但这个解法的缺点在于，为了在O(1)时间内执行搜索操作，每次插入一个新值都需要重新排列元素的位置。从而使得插入操作的解法平均时间复杂度变为O(N)。因此，算法总时间复杂度会变为O(N^2)。
鉴于插入操作的频繁，所以可以考虑使用二叉搜索树。
还需要在每个节点中放置一个计数器，以计算以此节点为根的子树中有多少个节点。
二叉搜索树实现
class TreeNode:
def __init__(self, val=0, cnt=0, left=None, right=None):
self.val = val
# 记录以此节点为根的子树中有多少个节点
self.cnt = cnt
self.left = left
self.right = right
class Solution:
# 建树函数
def insertIntoBST(self, root: Optional[TreeNode], val: int) -> Optional[TreeNode]:
if not root:
root = TreeNode(val, 1)
return root
root.cnt += 1
if val < root.val:
root.left = self.insertIntoBST(root.left, val)
else:
root.right = self.insertIntoBST(root.right, val)
return root
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.s = Solution()
self.tree = None
for num in nums:
self.tree = self.s.insertIntoBST(self.tree, num)
def searchK(self, root: Optional[TreeNode], kValue: int):
if not root:
return -1
# 获取右子树大小，避免访问 None.cnt报错
right_cnt = root.right.cnt if root.right else 0
# 找到当前 K 大元素
if kValue == right_cnt + 1:
return root.val
elif kValue <= right_cnt:
# K 在右子树中
return self.searchK(root.right, kValue)
else:
# K 在左子树中
return self.searchK(root.left, kValue - right_cnt - 1)
def add(self, val: int) -> int:
self.tree = self.s.insertIntoBST(self.tree, val)
return self.searchK(self.tree, self.k)- 此法的插入的平均复杂度是O(N)，提交后会出现时间超限。既如此，考虑使用优先队列。
优先队列实现
维护一个长度为K的优先队列，满队列后就弹出堆顶的最小元素（小顶堆），利用Python3内置的堆函数即可。
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.heap = []
for num in nums:
self.add(num)
def add(self, val: int) -> int:
# 利用heappush对插入元素进行局部排序
heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0]