使用Python寻找列表中最大、小的前K个元素及其索引_heapq提取前k个最大的元素

方法一：max()/min()函数的K次循环

无重复元素

K = 3
a = list(range(20))  # 无重复元素
Nmax_val_lis, Nmax_idx_lis = [], []
for _ in range(K):
    max_val = max(a)
    max_idx = a.index(max_val)

    Nmax_val_lis.append(max_val)
    Nmax_idx_lis.append(max_idx)

    a[max_idx] = float('-inf')

print(Nmax_val_lis)
print(Nmax_idx_lis)
1
2
3
4
5
6
7
8
9
10
11
12
13
14

    运行结果：

有重复元素

    取列表前K个最小元素也是同理，代码及运行结果如下：

K = 3
a = [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9]  # 存在重复元素
Nmin_val_lis, Nmin_idx_lis = [], []
for _ in range(K):
    min_val = min(a)
    min_idx = a.index(min_val)

    Nmin_val_lis.append(min_val)
    Nmin_idx_lis.append(min_idx)

    a[min_idx] = float('inf')

print(Nmin_val_lis)
print(Nmin_idx_lis)
1
2
3
4
5
6
7
8
9
10
11
12
13
14

    运行结果：

    从运行结果来看，K次循环的方法同样适用于存在重复元素的列表，所以该方法适用性很好。

方法二：Python中heapq包

无重复元素

import heapq

K = 3
a = list(range(10))  # 无重复元素
max_val_lis = heapq.nlargest(K, a)
max_idx_lis = list(map(a.index, max_val_lis))

print(max_val_lis)
print(max_idx_lis)
1
2
3
4
5
6
7
8
9

    运行结果：

    如果寻找列表的前K个最小值及索引也同理，使用heapq.nsmallest(K, a)。

有重复元素

import heapq

K = 3
a = [0, 1, 2, 3, 7, 7, 8, 8, 9, 9]
max_val_lis = heapq.nlargest(K, a)
max_idx_lis = list(map(a.index, max_val_lis))

print(max_val_lis)
print(max_idx_lis)
1
2
3
4
5
6
7
8
9

    运行结果：

    由于列表中最大值为9 9 8，对应索引分别是9 8 7/6，但是索引8重复输出了两次，索引9未被输出，这是因为list.index(val)方法只会输出list中val第一次出现的位置索引。因此有重复元素时，heapq.nlargest(K, a)方法不再适用！

    参考方法一，做如下修改：

import heapq

K = 3
a = [0, 1, 2, 3, 7, 7, 8, 8, 9, 9]
max_val_lis = heapq.nlargest(K, a)
max_idx_lis = []
for item in max_val_lis:
    idx = a.index(item)
    max_idx_lis.append(idx)
    a[idx] = float('-inf')

print(max_val_lis)
print(max_idx_lis)
1
2
3
4
5
6
7
8
9
10
11
12
13

    运行结果：

总结

1. 因为只需要求前K个大、小元素，Python也有对应的max()/min()函数包，因此只需要循环使用K次max()/min()即可（为了方便，暂且称这种方法为“K次循环法”）。该方法原理简单，且适用性极好，有无重复元素的列表均适用；
2. Python中有heapq.nlargest()和heapq.nsmallest()包方法，heapq包方法是基于堆的概念（放以后细讲）。若列表没有重复元素，该方法使用非常方便，一行代码即可；但若存在重复元素，该方法不能准确输出K个最大值对应的位置索引，因此该方法适用性一般。

 

 

    （本文完整的pdf请关注公众号“张张学算法”，并回复“016”获取~）