004找到最大或者最小的N个元素

heapq模块中的nlargest和nsmallest可以很方便的解决这个问题。

• nlargrst(int, [list|tuple|set|str])
• 返回值类型 list

import heapq
numbers = [1, 2, 4, 6, -4, -3, -9, 9, 2, 8]
large_nums = heapq.nlargest(3, numbers)
small_nums = heapq.nsmallest(3, numbers)
print(f'large nums: {large_nums}')
print(f'small nums: {small_nums}')

<OUTPUT>
large nums: [9, 8, 6]
small nums: [-9, -4, -3]

nlargest和nsmallest支持更加复杂的数据结构，接收参数key来声明将要处理的内容。

dict_in_list = [
    {'name': 'ws', 'age': 23},
    {'name': 'wfm', 'age': 31},
    {'name': 'scl', 'age': 28}
]
young = heapq.nlargest(2, dict_in_list, key=lambda s: s['age'])
elderly = heapq.nsmallest(2, dict_in_list, key=lambda s: s['age'])
print(f'young: {young}')
print(f'elderly: {elderly}')

<OUTPUT>
young: [{'name': 'wfm', 'age': 31}, {'name': 'scl', 'age': 28}]
elderly: [{'name': 'ws', 'age': 23}, {'name': 'scl', 'age': 28}]

可以使用heapify将列表转为堆，然后使用堆的特性来实现期望的功能；例如：堆的重要特性是heap[0]总是最小的元素。

heap = list(numbers)
# list -> heap(堆)
heapq.heapify(heap)print(f'heap: {heap}')

# 堆的重要特性是heap[0]总是最小的元素
# heapq.heappop()会将第一个元素（最小的）弹出，然后以第二小元素取而代之
print(heapq.heappop(heap))
print(heapq.heappop(heap))
print(heapq.heappop(heap))
print(heapq.heappop(heap))

<OUTPUT>
heap: [-9, -4, -3, 2, 2, 1, 4, 9, 6, 8]
-9
-4
-3
1

举例：找到最大的N个元素。

def search(lst: List, n: int):
    if n == 1:
        # 要找的元素只有一个时，max是最好的选择；当然，想要找最小的使用min
        return max(lst)
    elif n > (len(lst) * 0.9):
        # 要找的元素数量与序列长度相近时，先排序再切片效率会更高一些
        return sorted(lst)[:n]
    else:
        return heapq.nlargest(n, lst)