Python小顶堆的实现

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