堆的实现

最小的K个数。

堆插入、取出都是log(n)的复杂度。

class Heap:

def __init__(self):

self.data_list = []

def insert(self, x):

# 先把元素放在最后，然后从后往前依次堆化

# 这里以大顶堆为例，如果插入元素比父节点大，则交换，直到最后

self.data_list.append(x)

index = len(self.data_list) - 1

while index > 0:

p_index = index >> 1

if self.data_list[index] < self.data_list[p_index]:

self.data_list[index], self.data_list[p_index] = self.data_list[p_index], self.data_list[index]

index = p_index

def top(self):

return self.data_list[0]

def pop(self):

# 删除堆顶元素，然后将最后一个元素放在堆顶，再从上往下依次堆化

remove_data = self.data_list[0]

self.data_list[0] = self.data_list[-1]

self.data_list.pop()

self.heapify(0)

return remove_data

def heapify(self, index):

max_index = len(self.data_list) - 1

while index <= max_index:

maxvalue_index = index

left_index, right_index = (index << 1),  (index << 1) + 1

if (left_index <= max_index) and (self.data_list[left_index] < self.data_list[maxvalue_index]):

maxvalue_index = left_index

if (right_index <= max_index) and (self.data_list[right_index] < self.data_list[maxvalue_index]):

maxvalue_index = right_index

if maxvalue_index == index:

break

self.data_list[maxvalue_index], self.data_list[index] = self.data_list[index], self.data_list[maxvalue_index]

index = maxvalue_index

class Solution(object):

def smallestK(self, arr, k):

"""

:type arr: List[int]

:type k: int

:rtype: List[int]

"""

heapq = Heap()

for a in arr:

heapq.insert(a)

# print(heapq.data_list)

ans = []

for _ in range(k):

ans.append(heapq.pop())

return ans

python内置堆heapq

heap是一个list，具有堆特性的heap不一定是排好序的，使用heapq.heappop()出来的结果才是排好序的。

python默认的heapq模块只支持最小堆的构建。最大堆构建时，值可以乘-1，结果也再返回去。

heapq方法用法

heapq.heappush(heap, x)将x压入堆heap中

heapq.heappop(heap)从堆heap中弹出最小的元素

heapq.heapify(heap)让列表heap具备堆特征

heapq.heapreplace(heap, x)弹出最小的元素，并将x压入堆heap中

heapq.nlargest(n, iterable, key=None)返回iter中n个最大的元素

heapq.nsmallest(n, iterable, key=None)返回iter中n个最小的元素

example

import heapq

# 1.创建堆

# 方法一

nums = [3, 2, 4, 8, 3, 1, 4, 9, 4]

my_heapq_1 = []

for num in nums:

heapq.heappush(my_heapq_1, num)  # 加入堆

# 方法二

my_heapq_2 = [3, 2, 4, 8, 3, 1, 4, 9, 4]

heapq.heapify(my_heapq_2)

# 2. 访问堆顶

print([heapq.heappop(my_heapq_1) for _ in range(len(my_heapq_1))])  # 堆排序结果

print([heapq.heappop(my_heapq_2) for _ in range(len(my_heapq_2))])

# [1, 2, 3, 3, 4, 4, 4, 8, 9]

# 3.访问堆顶并替换

heapq.heapreplace(my_heapq_1, 23)

print([heapq.heappop(my_heapq_1) for _ in range(len(my_heapq_1))])  # 堆排序结果

# [2, 3, 3, 4, 4, 4, 8, 9, 23]

# 4.最大n个/最小n个。

nlargest = heapq.nlargest(5, my_heapq_1)

nsmallest = heapq.nsmallest(5, my_heapq_1)

print(nlargest, nsmallest)

# [9, 8, 4, 4, 4] [1, 2, 3, 3, 4]

元组堆

import heapq

# 1.创建堆

my_heapq = []

heapq.heappush(my_heapq, (5, 7))

heapq.heappush(my_heapq, (5, 4))

heapq.heappush(my_heapq, (7, 2))

heapq.heappush(my_heapq, (1, 3))

heapq.heappush(my_heapq, (3, 1))

print([heapq.heappop(my_heapq) for _ in range(len(my_heapq))])

# 2.最大n个/最小n个。

nlargest = heapq.nlargest(5, my_heapq, key=lambda e: (e[0], e[1]))

nsmallest = heapq.nsmallest(5, my_heapq, key=lambda e: (e[1], e[0]))

print(nlargest, nsmallest)