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"相同的函数,但我目前无法访问这些库.

Answer 1

在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 2

您可以创建一个带有可调用参数的函数,该参数将在列表推导的条件部分中使用.然后你可以使用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.

Answer 3

或者使用numpy的非零函数:

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

Answer 4

为什么不使用这个:

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

或者对于任意条件,f为您的条件定义一个函数并执行:

[i for i in range(len(a)) if f(a[i])]