сумма (условие) эквивалент в Python NumPy

Question

Я пытаюсь преобразовать часть кода Matlab в Python.

a=[1 2 3;4 5 6]
b= sum(a<5)
//output :
ans :
2 1 1

На самом деле возвращает количество элементов в каждом столбце, который имеет условие. Есть ли в numpy (python) эквивалентная функция для этого?

Answer 1

То же самое.

a=np.array([[1, 2, 3],[4, 5, 6]])
b=np.sum(a<5,axis=0) # the only difference is that you need to explicitly set the dimension

Answer 2

Хотя и не для этой цели, альтернативное решение будет

a=np.array([[1, 2, 3],[4, 5, 6]])
np.count_nonzero(a<5, axis=0)
# array([2, 1, 1])

Производительность

Для небольших массивов np.sum кажется немного быстрее

x = np.repeat([1, 2, 3], 100)
y = np.repeat([4, 5, 6], 100)
a=np.array([x,y])

%timeit np.sum(a<5, axis=0) 
# 7.18 µs ± 669 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit np.count_nonzero(a<5, axis=0)
# 11.8 µs ± 386 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Для очень больших массивов np.count_nonzero кажется немного быстрее

x = np.repeat([1, 2, 3], 5000000)
y = np.repeat([4, 5, 6], 5000000)
a=np.array([x,y])

%timeit np.sum(a<5, axis=0) 
# 126 ms ± 6.92 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.count_nonzero(a<5, axis=0)
# 100 ms ± 6.72 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)