数据结构之堆(heapq)

堆，或者说优先队列，是一颗完全二叉树

一般用数组存储，节点i的左、右子节点分别为2*i+1、2*i+2

堆分为小根堆和大根堆（python中为小根堆），小根堆的父节点的值小于或等于子节点的值

heapq

1. 创建堆

   import heapq
   
   heap=[]
   heapq.heapify(heap)

2. 添加元素

   heapq.heappush(heap,num)

3. 弹出元素

   弹出堆顶元素，对于python，是最小元素

   heapq.heappop(heap)

4. 取最值

   取n个最值，返回一个列表

   heapq.nlargest(2,heap)
   heapq.nsmallest(3,heap)

5. 堆合并

   heapq.merge(heap1,heap2)

6. 压入弹出

   注意分辨两者区别，尤其 num<heap[0] 的时候

   • heapq.pushpop(heap,num) # 先push num 再pop heap[0]
   • heapq.replace(heap,num) # 先pop heap[0] 再push num

   heapq.pushpop(heap,num)	# 先push再pop
   heapq.replace(heap,num) # 先pop再push

7. 一些注意点

   1. python中的heapq为小根堆，heap[0]为最小元素（C++中默认大根堆）
   2. python中可以将正整数堆取相反数，来达到大根堆的效果

数组中的第K个最大元素

直接调用堆就行了，虽然用排序也很方便

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heapq.heapify(nums)
        return heapq.nlargest(k,nums)[-1]

丑数 II

class Solution:
    def nthUglyNumber(self, n: int) -> int:
        heap=[1]
        seen={1}
        heapq.heapify(heap)
        for _ in range(n-1):
            tmp=heappop(heap)
            if 2*tmp not in seen:
                seen.add(2*tmp)
                heappush(heap,2*tmp)
            if 3*tmp not in seen:
                seen.add(3*tmp)
                heappush(heap,3*tmp)
            if 5*tmp not in seen:
                seen.add(5*tmp)
                heappush(heap,5*tmp)
        return heappop(heap)

相对名次

class Solution:
    def findRelativeRanks(self, score: List[int]) -> List[str]:
        n=len(score)
        prize=["Gold Medal", "Silver Medal", "Bronze Medal"]
        rank={i:(str(i) if i not in (1,2,3) else prize[i-1])  for i in range(1,n+1) }

        heapify(newscore:=deepcopy(score))
        dic={}
        for i in range(1,n+1):
            dic[heappop(newscore)]=i
        
        return [rank[n+1-dic[i]] for i in score]	# n+1-dic:因为python是小根堆