leetcode 703. 数据流中的第 K 大元素

背景

这题需要使用优先队列（小顶堆）来做，本质上是个完全二叉树，而对于一个完全二叉树来说，底层数据结构是可以用一个数组表示的。
对于任何一个节点arr[i]，左右子节点分别为arr[i*2+1]和arr[i*2+2]，其父节点为arr[ (i-1)/2 ]。
同时，小顶堆要求所有父节点都要小与子节点。

利用heapq库求解

由于Python只有hash map, 没有tree map,没法直接构建树结构的数据类型，所幸Python还有个heapq包可以用。

import heapq
class KthLargest:

    def __init__(self, k: int, nums):
        self.h = []
        self.k = k
        for item in nums:
            heapq.heappush(self.h, item)
            if len(self.h) > k:
                heapq.heappop(self.h)


    def add(self, val: int) -> int:
        heapq.heappush(self.h, val)
        if len(self.h) > self.k:
            heapq.heappop(self.h)
        return self.h[0]

手工实现小顶堆

如果想更深入理解小顶堆的底层实现机制可以手工打造一个。
首先定义一个小顶堆类Heap.

class Heap:
    def __init__(self):
        self.heap = []

    def push(self, val):
        self.heap.append(val)
        self._shift_up()

    def _swap(self, rsh, lsh):
        self.heap[rsh], self.heap[lsh] = self.heap[lsh], self.heap[rsh]

    def pop(self):
        item = self.heap[0]
        self._swap(0, -1)
        self.heap.pop()
        self._shift_down()
        return item

    def size(self):
        return len(self.heap)

    def top(self):
        return self.heap[0]

    def _shift_up(self):
        idx = len(self.heap) -1
        while idx > 0:
            parent = (idx - 1)//2
            if self.heap[parent] > self.heap[idx]:
                self._swap(parent, idx)
                idx = parent
            else:
                break 

    def _shift_down(self):
        idx = 0
        n = len(self.heap)
        while idx < n:
            left = idx *2 + 1
            right = idx * 2 + 2
            smaller = idx
            if left < n and  self.heap[left] < self.heap[smaller]:
                smaller = left
            if right < n and self.heap[right] < self.heap[smaller]:
                smaller = right

            if smaller == idx:
                break
            self._swap(smaller, idx)
            idx = smaller

当self.heap为空时，需要通过push方法将数据添加到heap的最后，由于每次添加了新数据，小顶堆的特性被破坏，需要上浮(shift_up)来更新self.heap，保证其依然是小顶堆结构。
当self.heap数据达到预期size，则需要弹出多余的数据，这时需要利用pop函数，注意小顶堆根节点保存的是整个堆的最小值，直接删除根节点是不可行的，这里的实现方法是将根节点(self.heap[0])与最右边的叶子节点(self.head[-1])进行交换，将要删除的数据移动到数组的最后，再用self.heap.pop()删除．