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

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