python - 根据条件获取数组索引位置

Question

我有一个大的 numpy 数组，我想根据给定的标准获取数组索引。Numpy 提供了这个但返回一个布尔数组：

>>> import numpy as np
>>> a = np.arrary([1, 2, 3, 4, 1, 2, 3]
>>> b = a == 3
>>> b
array([False, False, True, False, False, False, True])

但我真的很想将实际的索引位置作为整数，有没有比这更简单的方法：

>>> c = np.arange(len(b))
>>> c = c[b]
>>> c
array([2,6])

换句话说，有没有办法在不创建上面的 c 数组的情况下做到这一点？

Answer 1

我会去：

import numpy as np
a = np.array([1, 2, 3, 4, 1, 2, 3])
indices, = np.where(a==3)
print indices
# [2 6]

允许对 n-dim 数组进行更轻松的错误处理（即，如果要解包的项目过多，则会引发异常）并且不需要展平。

Answer 2

我相信您正在寻找numpy.argwhere：

In[1]: import numpy as np
In[2]: a = np.array([1,2,3,4,1,2,3])
In[3]: b = np.argwhere(a==3).flatten()
In[4]: b
Out[1]: array([2, 6])