class TopkHeap(object):
def __init__(self, k):
self.k = k
self.data = []
def Push(self, elem):
if len(self.data) < self.k:
heapq.heappush(self.data, elem)
else:
topk_small = self.data[0]
if elem > topk_small:
heapq.heapreplace(self.data, elem)
def TopK(self):
#return [x for x in reversed([heapq.heappop(self.data) for _ in range(len(self.data))])]
return reversed(sorted(self.data))
自己实现小顶堆
class MinHeap:
def __init__(self, size):
self.size = size
self.data = []
def pushpop(self, v):
if len(self.data) < self.size:
self.data.append(v)
self.swim(len(self.data) - 1)
elif v > self.data[0]:
self.data[0] = v
self.sink(0)
def large(self, p, q):
return self.data[p] > self.data[q]
def exch(self, p, q):
self.data[p], self.data[q] = self.data[q], self.data[p]
def swim(self, node):
while (node - 1) // 2 >= 0 and self.large((node - 1) // 2, node):
self.exch((node - 1) // 2, node)
node = (node - 1) // 2
def sink(self, node):
while node * 2 < self.size - 1:
left = node * 2 + 1
right = node * 2 + 2
child_min = left
if right < self.size and self.large(left, right):
child_min = right
if not self.large(node, child_min):
break
self.exch(node, child_min)
node = child_min
class Solution:
def getLargestNumbers(self, arr: List[int], k: int) -> List[int]:
if not k:
return []
mq = MinHeap(k)
for num in arr:
mq.pushpop(num)
return sorted(mq.data, reverse=True)
变态的需求来了:给出N长的序列,求出BtmK小的元素,即使用大顶堆。
概括一种最简单的:
将push(e)改为push(-e)、pop(e)改为-pop(e)。
也就是说,在存入堆、从堆中取出的时候,都用相反数,而其他逻辑与TopK完全相同,看代码:
class BtmkHeap(object):
def __init__(self, k):
self.k = k
self.data = []
def Push(self, elem):
# Reverse elem to convert to max-heap
elem = -elem
# Using heap algorighem
if len(self.data) < self.k:
heapq.heappush(self.data, elem)
else:
topk_small = self.data[0]
if elem > topk_small:
heapq.heapreplace(self.data, elem)
def BtmK(self):
return sorted([-x for x in self.data])
自己实现大顶堆
class MaxHeap:
def __init__(self, size):
self.size = size
self.data = []
def pushpop(self, v):
if len(self.data) < self.size:
self.data.append(v)
self.swim(len(self.data) - 1)
elif v < self.data[0]:
self.data[0] = v
self.sink(0)
def less(self, p, q):
return self.data[p] <= self.data[q]
def exch(self, p, q):
self.data[p], self.data[q] = self.data[q], self.data[p]
def swim(self, node):
while (node - 1) // 2 >= 0 and self.less((node - 1) // 2, node):
self.exch((node - 1) // 2, node)
node = (node - 1) // 2
def sink(self, node):
while node * 2 < self.size - 1:
left = node * 2 + 1
right = node * 2 + 2
child_max = left
if right < self.size and self.less(left, right):
child_max = right
if not self.less(node, child_max):
break
self.exch(node, child_max)
node = child_max
class Solution:
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
if not k:
return []
mq = MaxHeap(k)
for num in arr:
mq.pushpop(num)
return sorted(mq.data)