使用Python列表推导式基于条件查找元素的索引

Question

对于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]

- Lee

---

那么 [idx for idx in range(len(a)) if a[idx] > 2] 怎么样？这在 Python 中有点棘手，因为它不像其他语言那样经常使用索引。 - NullUserException

Answer 1

另一种方法：

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

总的来说，要记住虽然find是一个现成的函数，但列表推导式是一般化的，因此非常强大的解决方案。没有什么会阻止你在Python中编写一个find函数，并随后根据需要使用它。例如：

>>> 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。更符合Python风格的做法是使用生成器和迭代器。

- Eli Bendersky

---

OP本可以将其写为[i for i,v in enumerate(a) if v > 2]。 - NullUserException

这并不更短，反而更长了。在原始代码中将“index”替换为“i”，将“value”替换为“v”，然后计算字符数。 - agf

@NullUser，agf：您说得对，但重点在于第二部分 :) - Eli Bendersky

使用 enumerate 覆盖 range(len(...)) 既更健壮又更高效。 - Mike Graham

@Mike Graham：我同意 - 将更改find_indices函数以使用enumerate。 - Eli Bendersky

Answer 2

• 在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])

- Mike Graham

---

你有两个列表，一个是范围值，一个是角度值，你想要过滤掉超过某个阈值的范围值。如何以“最佳方式”过滤掉相应的角度值？ - Mehdi

filtered_ranges_and_angles = [(range, angle) for range, angle in zip(ranges, angles) if should_be_kept(range)] 的意思是：根据 should_be_kept(range) 的结果过滤出符合条件的测量范围和角度，并将它们以元组的形式存储在 filtered_ranges_and_angles 变量中。 - Mike Graham

在Python中，你根本不会使用索引来处理这个问题，而是直接处理值。这个语句表明你还没有做足数据分析和机器学习建模。基于某些条件的一个张量的索引被用来过滤另一个张量。 - horaceT

Answer 3

对我来说它很有效：

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

- Alexander Cyberman

Answer 4

也许另外一个问题是，“当你得到这些索引后，你会用它们做什么？” 如果你要用它们来创建另一个列表，在Python中，它们是一个不必要的中间步骤。如果你想要所有与给定条件匹配的值，只需使用内置过滤器：

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

或者编写自己的列表推导式：

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

如果您想将它们从列表中移除，则Pythonic的方法不是直接从列表中删除，而是编写一个列表理解（list comprehension），如同您在创建新列表一样，然后使用左侧的listvar[:]将其重新分配到原位置：

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

Matlab提供find功能，因为它的基于数组的模型通过使用数组索引来选择项目。你当然可以在Python中这样做，但更符合Pythonic的方式是使用迭代器和生成器，正如@EliBendersky已经提到的。

- PaulMcG

---

Paul，我在脚本/函数/类中还没有遇到过这种需要。这更多是为了交互式测试我正在编写的类。 - Lee

@Mike - 感谢您的编辑，但我真的是想说 a[:] = ... - 请参考Alex Martelli在此问题https://dev59.com/7nM_5IYBdhLWcg3wgjS2中的回答。 - PaulMcG

@Paul，我假设（也希望如此！）你在描述中所说的“创建新列表”并不是真的；我发现当程序尽可能地少改变现有数据时，它们更容易理解和维护。无论如何，我很抱歉越界了——你肯定可以编辑你的帖子，恢复到任何你想要的状态。 - Mike Graham

Answer 5

即使这是一个晚答案：我认为这仍然是一个非常好的问题，而且在我看来，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]

在我的想象中，最完美的方法是创建一个继承自列表的子类，并将索引函数添加为类方法。这样只需要使用过滤方法：

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

- Gerhard Hagerer

Answer 6

The following should then work (but I haven't got a Python interpreter on hand to test it):

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]

这是一个不错的解决方案。但内置类型并不适合被子类化。您可以使用组合而不是继承。这应该可以解决问题：

import numpy

class my_array:
    def __init__(self, data):
        self.data = numpy.array(data)

    def find(self, b):
        r = numpy.array(list(range(len(self.data))))
        return list(r[b])

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

- Arthur Bulcão Gomes