手撕
def swim(heap, i):
'''上浮, 插入一个新元素后,尽量往上爬'''
while i > 0:
pi = (i - 1) // 2
if heap[pi] > heap[i]:
heap[pi], heap[i] = heap[i], heap[pi]
i = pi
else:
break
def sink(heap, i, n):
'''如果要pop堆顶元素,可以将堆顶元素与堆尾交换,然后堆顶下沉'''
while i < n:
l, r = i * 2 + 1, i * 2 + 2
if l >= n:
break
min_i = l
if r < n and heap[r] < heap[l]:
min_i = r
if heap[min_i] < heap[i]:
heap[min_i], heap[i] = heap[i], heap[min_i]
i = min_i
else:
break
def push(heap, num, k):
if len(heap) == k:
if num > heap[0]:
heap[0] = num
sink(heap, 0, k)
else:
heap.append(num)
swim(heap, len(heap) - 1)
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.nums = nums
self.heap = []
for num in nums:
push(self.heap, num, k)
def add(self, val: int) -> int:
push(self.heap, val, self.k)
return self.heap[0]
调包
def push(heap, num, k):
if len(heap) == k:
if num > heap[0]:
heap[0]=num
heapq.heapify(heap)
else:
heapq.heappush(heap, num)
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.nums = nums
self.heap = []
for num in nums:
push(self.heap, num, k)
def add(self, val: int) -> int:
push(self.heap, val, self.k)
return self.heap[0]