python取列表中前n个最大、最小值,及对应的下标

使用heapq这个库可以帮助我们实现这样的功能,通过构造优先队列,得到想要的结果,这里封装以下变为功能函数:

def get_larger_from_list(my_list: list, n=1):
    """获取列表中最大的max_num个数,以及对应的下标"""
    my_list = copy.deepcopy(my_list)
    larger_number: list = heapq.nlargest(n, my_list)
    larger_index = []
    for _ in larger_number:
        index = my_list.index(_)  # 获取最小值对应的第一个下标
        larger_index.append(index)
        my_list[index] = -float("inf")
    return larger_number, larger_index


def get_smaller_from_list(my_list: list, n=1):
    """获取列表中最小的min_num个数,以及对应的下标"""
    my_list = copy.deepcopy(my_list)
    smaller_number: list = heapq.nsmallest(n, my_list)
    smaller_index = []
    for _ in smaller_number:
        index = my_list.index(_)  # 获取最小值对应的第一个下标
        smaller_index.append(index)
        my_list[index] = float("inf")
    return smaller_number, smaller_index

使用循环依次取数,可以确保如果有多个重复值时,同样可以返回正确的结果

示例代码

import copy
import heapq
import numpy as np


def get_larger_from_list(my_list: list, n=1):
    """获取列表中最大的max_num个数,以及对应的下标"""
    my_list = copy.deepcopy(my_list)
    larger_number: list = heapq.nlargest(n, my_list)
    larger_index = []
    for _ in larger_number:
        index = my_list.index(_)  # 获取最小值对应的第一个下标
        larger_index.append(index)
        my_list[index] = -float("inf")
    return larger_number, larger_index


def get_smaller_from_list(my_list: list, n=1):
    """获取列表中最小的min_num个数,以及对应的下标"""
    my_list = copy.deepcopy(my_list)
    smaller_number: list = heapq.nsmallest(n, my_list)
    smaller_index = []
    for _ in smaller_number:
        index = my_list.index(_)  # 获取最小值对应的第一个下标
        smaller_index.append(index)
        my_list[index] = float("inf")
    return smaller_number, smaller_index


def main():
    random_list = list(np.random.randint(0, 100, size=100))
    smaller_value, smaller_index = get_smaller_from_list(random_list, 10)
    larger_value, larger_index = get_larger_from_list(random_list, 10)
    print('最小值:{}\n对应下标:{}'.format(smaller_value, smaller_index))
    print('最大值:{}\n对应下标:{}'.format(larger_value, larger_index))


if __name__ == '__main__':
    main()

得到结果:

最小值:[0, 0, 0, 1, 1, 2, 4, 4, 6, 6]
对应下标:[9, 58, 67, 5, 21, 66, 1, 2, 28, 40]
最大值:[99, 99, 98, 97, 94, 93, 90, 88, 88, 87]
对应下标:[17, 88, 12, 65, 36, 84, 25, 39, 60, 15]

猜你喜欢

转载自blog.csdn.net/weixin_35757704/article/details/129163531