Получить позиции индекса массива на основе условного

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] 

Позволяет для облегчения обработки ошибок при п-тусклых массивов (то есть, исключение будет если слишком много предметов для распаковки), и это не требует сглаживания.

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