使用Python列表推导根据条件查找元素索引

Question

以下的Python代码在从Matlab转过来的时候，看起来有点啰嗦。

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

Answer 1

对我来说，这个方法很好用：

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

Answer 2

• 在Python中，你根本不需要用索引来处理这些事情，而是直接使用值——[value for value in a if value > 2]。通常，如果你在用索引，那就说明你可能没有用最好的方法。

• 如果你确实需要一个和Matlab类似的接口，可以使用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])
  

Answer 3

还有一种方法：

>>> [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的风格。