堆的实现
最小的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)