使用python list comprehension根据条件查找元素的索引


119

当来自Matlab背景时,以下Python代码似乎很冗长

>>> a = [1, 2, 3, 1, 2, 3]
>>> [index for index,value in enumerate(a) if value > 2]
[2, 5]

在Matlab中,我可以写:

>> a = [1, 2, 3, 1, 2, 3];
>> find(a>2)
ans =
     3     6

是否有使用Python编写此代码的简便方法,还是只保留长版本?


感谢您对Python语法原理的所有建议和解释。

在numpy网站上找到以下内容后,我想我已经找到了喜欢的解决方案:

http://docs.scipy.org/doc/numpy/user/basics.indexing.html#boolean-or-mask-index-arrays

将来自该网站的信息应用于上述我的问题,将得到以下结果:

>>> from numpy import array
>>> a = array([1, 2, 3, 1, 2, 3])
>>> b = a>2 
array([False, False, True, False, False, True], dtype=bool)
>>> r = array(range(len(b)))
>>> r(b)
[2, 5]

然后,下面的内容应该可以工作(但是我手头没有Python解释器来对其进行测试):

class my_array(numpy.array):
    def find(self, b):
        r = array(range(len(b)))
        return r(b)


>>> a = my_array([1, 2, 3, 1, 2, 3])
>>> a.find(a>2)
[2, 5]

6
怎么[idx for idx in range(len(a)) if a[idx] > 2]样 在Python中这样做有点尴尬的原因是,它使用的索引不如其他语言那么多。
NullUserException 2011年

Answers:


77
  • 在Python中,您根本不会为此使用索引,而只处理值[value for value in a if value > 2]。通常,处理索引意味着您没有采取最佳方法。

  • 如果确实需要类似于Matlab的API,则可以使用numpy,这是Python中用于多维数组和数值数学的软件包,受Matlab的启发很大。您将使用numpy数组而不是列表。

    >>> import numpy
    >>> a = numpy.array([1, 2, 3, 1, 2, 3])
    >>> a
    array([1, 2, 3, 1, 2, 3])
    >>> numpy.where(a > 2)
    (array([2, 5]),)
    >>> a > 2
    array([False, False,  True, False, False,  True], dtype=bool)
    >>> a[numpy.where(a > 2)]
    array([3, 3])
    >>> a[a > 2]
    array([3, 3])

2
您有一个列表,一个用于范围,一个用于角度,您要过滤掉超出某个阈值的范围值。您还如何以“最佳方式”过滤与这些范围相对应的角度?
Mehdi

3
filtered_ranges_and_angles = [(range, angle) for range, angle in zip(ranges, angles) if should_be_kept(range)]
Mike Graham

7
“在Python中,您根本不会使用索引,而只处理值”。该语句表明您尚未完成足够的数据分析和机器学习建模。一个基于特定条件的张量的索引用于过滤另一个张量。
horaceT

63

其他方式:

>>> [i for i in range(len(a)) if a[i] > 2]
[2, 5]

通常,请记住,虽然这find是一个现成的函数,但列表推导是一个通用的解决方案,因此非常有效。没有什么可以阻止您find使用Python 编写函数并在以后根据需要使用它。即:

>>> def find_indices(lst, condition):
...   return [i for i, elem in enumerate(lst) if condition(elem)]
... 
>>> find_indices(a, lambda e: e > 2)
[2, 5]

请注意,我在这里使用列表来模仿Matlab。使用生成器和迭代器会更Pythonic。


2
OP可以改写它[i for i,v in enumerate(a) if v > 2]
NullUserException 2011年

这不是更短,而是更长。更换indexivaluev原和计数的字符。
2011年

@ NullUser,AGF:您是对的,但要点是第二部分:)
Eli Bendersky

1
使用enumerateover range(len(...))既健壮又高效。
Mike Graham

1
@迈克·格雷厄姆(Mike Graham):我同意-将更改find_indices使用的功能enumerate
伊莱·

22

对我来说,效果很好:

>>> import numpy as np
>>> a = np.array([1, 2, 3, 1, 2, 3])
>>> np.where(a > 2)[0]
[2 5]

6

也许另一个问题是,“一旦获得这些索引,您将如何处理这些索引?” 如果要使用它们创建另一个列表,那么在Python中,它们是不必要的中间步骤。如果想要所有与给定条件匹配的值,只需使用内置过滤器:

matchingVals = filter(lambda x : x>2, a)

或编写您自己的列表理解:

matchingVals = [x for x in a if x > 2]

如果要从列表中删除它们,那么Python的方法不一定是从列表中删除,而是像编写新列表一样编写列表理解,然后使用listvar[:]左侧的就地分配-侧:

a[:] = [x for x in a if x <= 2]

Matlab find之所以提供它,是因为其以数组为中心的模型通过使用数组索引选择项目而起作用。当然,您可以在Python中执行此操作,但更Pythonic的方式是使用迭代器和生成器,如@EliBendersky所述。


保罗,我还没有在脚本/函数/类中遇到过这种需要。它更多地用于我正在编写的类的交互式测试。
李·

@Mike-感谢您的编辑,但我确实的意思是a[:] = ...-请参阅Alex Martelli对这个问题的回答stackoverflow.com/questions/1352885/…
PaulMcG 2011年

@Paul,我假设(并希望!)从描述中您并不是真的要说要“创建新列表”;我发现程序在非常少量地更改现有数据时倾向于更易于理解和维护。无论如何,我很抱歉超载-您当然应该能够将您的帖子修改回您想要的任何内容。
Mike Graham

6

即使答案很晚:我认为这仍然是一个很好的问题,而且恕我直言,Python(没有其他库或工具包(例如numpy))仍然缺乏方便的方法来根据手动定义的过滤器访问列表元素的索引。

您可以手动定义一个提供该功能的功能:

def indices(list, filtr=lambda x: bool(x)):
    return [i for i,x in enumerate(list) if filtr(x)]

print(indices([1,0,3,5,1], lambda x: x==1))

产量:[0,4]

在我的想象中,完美的方法将是创建列表的子类并添加索引作为类方法。这样,只需要使用filter方法:

class MyList(list):
    def __init__(self, *args):
        list.__init__(self, *args)
    def indices(self, filtr=lambda x: bool(x)):
        return [i for i,x in enumerate(self) if filtr(x)]

my_list = MyList([1,0,3,5,1])
my_list.indices(lambda x: x==1)

我在这里详细介绍了该主题:http//tinyurl.com/jajrr87

By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.