numpy: подсчет в 2D-массиве, где элемент удовлетворяет условию

Question

У меня есть массив numpy, как показано ниже. Мне нужно количество строк, где первый элемент равен 2. Итак, в приведенном ниже массиве четыре строки начинаются с 2 - ответ будет 4. Как это лучше всего сделать в numpy? (Я не могу использовать панды, но могу использовать scipy).

array([[1, 4, 5], 
     [1, 4, 5], 
     [2, 4, 5], 
     [2, 4, 5], 
     [2, 4, 5], 
     [2, 4, 5], 
     [3, 4, 5], 
     [3, 4, 5], 
     [3, 4, 5], 
     [3, 4, 5], 
     [3, 4, 5], 
     [3, 4, 5]]) 

Answer 1

Во-первых, сделать первый столбец, все строки:

a[:,0] 

Затем найдите 2 s:

a[:,0] == 2 

Это дает булево массив. Что вы можете суммировать:

(a[:,0] == 2).sum() 

Answer 2

Существует np.count_nonzero, которые в общих идиомах применяются для логических массивов, полученных путем оценки условия

np.count_nonzero(data[:, 0] == 2) 

Btw. это, вероятно, только ради примера, но если массив отсортирован, как ваши вы можете также использовать np.searchsorted

np.diff(np.searchsorted(data[:, 0], (2, 3)))[0] 

Answer 3

Еще один подход в дополнение к выше подходов

>>> x[:,0]==2 
array([False, False, True, True, True, True, False, False, False, 
     False, False, False], dtype=bool) 

даст вам значение TRUE для строк, которые имеют первую колонку 2.

>>> x[x[:,0]==2] 
array([[2, 4, 5], 
     [2, 4, 5], 
     [2, 4, 5], 
     [2, 4, 5]]) 

дает вам соответствуют строки и удовлетворяют требуемому условию. Теперь вы можете использовать функцию формы для получения длины.

x[x[:,0]==2].shape[0]