Python小顶堆的实现
以前都是用的heapq,这次面试碰到让自己实现一个堆得到top k的元素。抛砖引玉,希望各位提些意见指正
class Heap:
def __init__(self, k):
self.val = []
self.k = k
def get_left(self, idx):
return idx * 2 + 1 if len(self.val) > idx * 2 + 1 else None
def get_right(self, idx):
return idx * 2 + 2 if len(self.val) > idx * 2 + 2 else None
def get_parent(self, idx):
return (idx - 1) // 2 if idx > 0 else None
def sift_up(self, idx):
cur = idx
parent = self.get_parent(idx)
while parent is not None and self.val[cur] < self.val[parent]:
# print(parent, cur, self.val[parent], self.val[cur])
temp = self.val[cur]
self.val[cur] = self.val[parent]
self.val[parent] = temp
cur = parent
parent = self.get_parent(cur)
def insert(self, n):
if len(self.val) < self.k:
self.val.append(n)
idx = len(self.val) - 1
self.sift_up(idx)
else:
if n > self.val[0]:
self.delete()
self.val.append(n)
idx = len(self.val) - 1
self.sift_up(idx)
return
def sift_down(self, idx):
l = self.get_left(idx)
r = self.get_right(idx)
cur = idx
if not l:
min_child = r
elif not r:
min_child = l
else:
min_child = l if self.val[l] == min(self.val[l], self.val[r]) else r
while cur is not None and min_child is not None and self.val[cur] > self.val[min_child]:
# print(cur, min_child, self.val[cur], self.val[min_child])
temp = self.val[cur]
self.val[cur] = self.val[min_child]
self.val[min_child] = temp
cur = min_child
l = self.get_left(cur)
r = self.get_right(cur)
if not l:
min_child = r
elif not r:
min_child = l
else:
min_child = l if self.val[l] == min(self.val[l], self.val[r]) else r
def delete(self):
if self.val:
temp = self.val[0]
self.val[0] = self.val[-1]
self.val[-1] = temp
n = self.val.pop()
self.sift_down(0)
return n
return None