python手写实现堆数据结构

1.初始化

def __init__(self,desc=False):
    """
    初始化，默认创建一个小顶堆
    """
    self.heap = []
    self.desc = desc

2.堆的大小

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

3.返回堆顶元素

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

4.添加元素

def push(self,item):
    """
    添加元素
    第一步，把元素加入到数组末尾
    第二步，把末尾元素向上调整
    """
    self.heap.append(item)
    self._sift_up(self.size-1)

5.弹出堆顶元素

def pop(self):
    """
    弹出堆顶
    第一步，记录堆顶元素的值
    第二步，交换堆顶元素与末尾元素
    第三步，删除数组末尾元素
    第四步，新的堆顶元素向下调整
    第五步，返回答案
    """
    item = self.heap[0]
    self._swap(0,self.size-1)
    self.heap.pop()
    self._sift_down(0)
    return item

6.判断两个元素的大小关系，这里有个小trick

return left if True else right #条件为True则总是返回left
return left if False else right #条件为False则总是返回right

def _smaller(self,lhs,rhs):
    return lhs > rhs if self.desc else lhs < rhs

7.向上调整

def _sift_up(self,index):
    """
    向上调整
    如果父节点和当前节点满足交换的关系
    （对于小顶堆是父节点元素更大，对于大顶堆是父节点更小），
    则持续将当前节点向上调整
    """
    while index:
        parent = (index-1) // 2
        
        if self._smaller(self.heap[parent],self.heap[index]):
            break
            
        self._swap(parent,index)
        index = parent

8.向下调整

def _sift_down(self,index):
    """
    向下调整
    如果子节点和当前节点满足交换的关系
    （对于小顶堆是子节点元素更小，对于大顶堆是子节点更大），
    则持续将当前节点向下调整
    """
    # 若存在子节点
    while index*2+1 < self.size:
        smallest = index
        left = index*2+1
        right = index*2+2
        
        if self._smaller(self.heap[left],self.heap[smallest]):
            smallest = left
            
        if right < self.size and self._smaller(self.heap[right],self.heap[smallest]):
            smallest = right
            
        if smallest == index:
            break

        self._swap(index,smallest)
        index = smallest

9.交换两个元素

def _swap(self,i,j):
    self.heap[i],self.heap[j] = self.heap[j],self.heap[i]

完整代码

class Heap:
    def __init__(self,desc=False):
        """
        初始化，默认创建一个小顶堆
        """
        self.heap = []
        self.desc = desc
    
    @property
    def size(self):
        return len(self.heap)
    
    def top(self):
        if self.size:
            return self.heap[0]
        return None
    
    def push(self,item):
        """
        添加元素
        第一步，把元素加入到数组末尾
        第二步，把末尾元素向上调整
        """
        self.heap.append(item)
        self._sift_up(self.size-1)
    
    def pop(self):
        """
        弹出堆顶
        第一步，记录堆顶元素的值
        第二步，交换堆顶元素与末尾元素
        第三步，删除数组末尾元素
        第四步，新的堆顶元素向下调整
        第五步，返回答案
        """
        item = self.heap[0]
        self._swap(0,self.size-1)
        self.heap.pop()
        self._sift_down(0)
        return item
    
    def _smaller(self,lhs,rhs):
        return lhs > rhs if self.desc else lhs < rhs
    
    def _sift_up(self,index):
        """
        向上调整
        如果父节点和当前节点满足交换的关系
        （对于小顶堆是父节点元素更大，对于大顶堆是父节点更小），
        则持续将当前节点向上调整
        """
        while index:
            parent = (index-1) // 2
            
            if self._smaller(self.heap[parent],self.heap[index]):
                break
                
            self._swap(parent,index)
            index = parent
    
    def _sift_down(self,index):
        """
        向下调整
        如果子节点和当前节点满足交换的关系
        （对于小顶堆是子节点元素更小，对于大顶堆是子节点更大），
        则持续将当前节点向下调整
        """
        # 若存在子节点
        while index*2+1 < self.size:
            smallest = index
            left = index*2+1
            right = index*2+2
            
            if self._smaller(self.heap[left],self.heap[smallest]):
                smallest = left
                
            if right < self.size and self._smaller(self.heap[right],self.heap[smallest]):
                smallest = right
                
            if smallest == index:
                break
            
            self._swap(index,smallest)
            index = smallest
    
    def _swap(self,i,j):
        self.heap[i],self.heap[j] = self.heap[j],self.heap[i]

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

题解：

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.heap = Heap()
        self.k = k
        for num in nums:
            self.heap.push(num)
            if self.heap.size > k:
                self.heap.pop()


    def add(self, val: int) -> int:
        self.heap.push(val)
        if self.heap.size > self.k:
            self.heap.pop()
        return self.heap.top()

作者：MiloMusiala
链接：https://leetcode-cn.com/problems/kth-largest-element-in-a-stream/solution/python-dong-hua-shou-xie-shi-xian-dui-by-ypz2/
来源：力扣（LeetCode）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。