使用Python寻找列表中最大、小的前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)

    运行结果:

image-20230310111644056

有重复元素

    取列表前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)

    运行结果:

image-20230310111608904

    从运行结果来看,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)

    运行结果:

image-20230310112518368

    如果寻找列表的前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)

    运行结果:

image-20230310132634414

    由于列表中最大值为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)

    运行结果:

image-20230310133526512

总结

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

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

猜你喜欢

转载自blog.csdn.net/weixin_40583722/article/details/129441680