Python中类似MATLAB的find()函数

Question

在MATLAB中，找到满足特定条件的值的索引非常简单：

>> a = [1,2,3,1,2,3,1,2,3];
>> find(a > 2)     % find the indecies where this condition is true
[3, 6, 9]          % (MATLAB uses 1-based indexing)
>> a(find(a > 2))  % get the values at those locations
[3, 3, 3]

那么在Python中，最好的方法是什么呢？

到目前为止，我想出了以下方法。只是想获取这些值：

>>> a = [1,2,3,1,2,3,1,2,3]
>>> [val for val in a if val > 2]
[3, 3, 3]

但是如果我想要这些值的索引，那就有点复杂了：

>>> a = [1,2,3,1,2,3,1,2,3]
>>> inds = [i for (i, val) in enumerate(a) if val > 2]
>>> inds
[2, 5, 8]
>>> [val for (i, val) in enumerate(a) if i in inds]
[3, 3, 3]

在Python中有没有更好的方法，特别是对于任意条件（不仅仅是'val > 2'）？

我在NumPy中找到了与MATLAB的'find'函数等效的函数，但我目前无法使用这些库。

相关问题：

• 满足值和索引条件的NumPy数组元素索引
• numpy数组 - 替换特定值
• 在Python/Numpy中选择满足条件的特定行，且只选择某些列
• 将Matlab中的1:n数组转换为Python
• 如何切片numpy数组以获取前两行和后两行？
• 在Python中，有简单的方法根据值条件构建子列表吗？
• 从numpy数组提取数据
• 理解numpy.r_()拼接的语法

Answer 1

或者使用numpy的nonzero函数：

import numpy as np
a    = np.array([1,2,3,4,5])
inds = np.nonzero(a>2)
a[inds] 
array([3, 4, 5])

Answer 2

在numpy中，你可以使用 where :

>> import numpy as np
>> x = np.random.randint(0, 20, 10)
>> x
array([14, 13,  1, 15,  8,  0, 17, 11, 19, 13])
>> np.where(x > 10)
(array([0, 1, 3, 6, 7, 8, 9], dtype=int64),)

Answer 3

你可以写一个函数，这个函数可以接收一个可调用的参数，这个参数会在你列表推导式的条件部分使用。然后你可以用一个lambda表达式或者其他的函数对象来传递你想要的条件：

def indices(a, func):
    return [i for (i, val) in enumerate(a) if func(val)]

a = [1, 2, 3, 1, 2, 3, 1, 2, 3]

inds = indices(a, lambda x: x > 2)

>>> inds
[2, 5, 8]

这样做就更接近你在Matlab中的例子，而且不需要加载整个numpy库。