在Python中获取列表中最小的N个元素的索引

3
我想要获取列表中最小的N个元素的索引。如果我能够将输出放在另一个列表中,那就太好了。
例如:
[1, 1, 10, 5, 3, 5]
output = [0, 1]

[10, 5, 12, 5, 0, 10]
output = [4]

[9, 2, 8, 2, 3, 4, 2]
output = [1, 3, 6]

[10, 10, 10, 10, 10, 10]
output = [0, 1, 2, 3, 4, 5]

我知道.index返回列表中最小值的第一个索引,但当最小值出现多次时,我不知道如何返回所有最小值的索引。

4
不接受我的答案,接受Wim的答案... - Avinash Raj
1
这都是关于性能的问题。但我的也能正常工作。 - Avinash Raj
2个回答

7
>>> L = [9, 2, 8, 2, 3, 4, 2]
>>> minL = min(L)
>>> [i for i, x in enumerate(L) if x == minL]
[1, 3, 6]

目前,其他解决方案在迭代过程中会调用min,导致复杂度为O(n^2)


编辑说明:天真的解决方案的n^2复杂度证明:

>>> L1000 = [randint(0, 100) for _ in xrange(1000)]
>>> L2000 = [randint(0, 100) for _ in xrange(2000)]
>>> L3000 = [randint(0, 100) for _ in xrange(3000)]
>>> L4000 = [randint(0, 100) for _ in xrange(4000)]
>>> L5000 = [randint(0, 100) for _ in xrange(5000)]
>>> timeit [i for i, x in enumerate(L1000) if x == min(L1000)]
10 loops, best of 3: 18.8 ms per loop
>>> timeit [i for i, x in enumerate(L2000) if x == min(L2000)]
10 loops, best of 3: 73.6 ms per loop
>>> timeit [i for i, x in enumerate(L3000) if x == min(L3000)]
1 loops, best of 3: 166 ms per loop
>>> timeit [i for i, x in enumerate(L4000) if x == min(L4000)]
1 loops, best of 3: 294 ms per loop
>>> timeit [i for i, x in enumerate(L5000) if x == min(L5000)]
1 loops, best of 3: 457 ms per loop

在每次迭代中,if x == min(L) 是否被调用了? - itzMEonTV
这对于标准库的答案来说要快得多,+1。如果您有兴趣,我在我的答案中添加了一个大列表基准测试 :) - miradulo

4

您也可以使用NumPy.where一起使用,这对于更大的列表(更多元素意味着数千个或更多)非常适用。

import numpy as np
mylist = [1, 1, 10, 5, 1, 5]
minL = min(mylist)
numparray = np.array(mylist)
print(list( np.where(numparray == minL)[0]))

输出:

[0, 1, 4]

简单基准测试:

Wim的回答:

>>>setup = '''
L = [9, 2, 8, 2, 3, 4, 2] * 5000
minL = min(L)
'''
>>> print (min(timeit.Timer(
            '[i for i, x in enumerate(L) if x == minL]', setup=setup).repeat(7, 1000)))
3.218847516924143

唐奇·康的回答:

>>> setup = '''
import numpy as np
mylist = [9, 2, 8, 2, 3, 4, 2] * 5000
minL = min(mylist)
numparray = np.array(mylist)
'''
>>> print (min(timeit.Timer(
          'list( np.where(numparray == minL)[0])', setup=setup).repeat(7, 1000)))
1.4015787709504366

对于包含35000个元素的列表,NumPy解决方案的速度是Wim答案的两倍以上。然而对于小型列表,Wim的答案已经足够。
此外,一个小型基准测试可以展示在列表推导式中调用min函数的O(n^2)时间复杂度(700个元素的列表)。
>>> setup = '''
L = [9, 2, 8, 2, 3, 4, 2] * 100
'''
>>> print (min(timeit.Timer(
           '[i for i,j in enumerate(L) if j == min(L)]', setup=setup).repeat(7, 1000))) 
7.43799185753

4
如果您可使用NumPy,那么它是实现此操作的最佳方式。 - wim
你也可以尝试使用numpy.min代替min,但在我的机器上差异很小。 - dshepherd
@dshepherd 是的,我认为效率差异可以忽略不计,因为我已经有一个内置列表了。 - miradulo

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接