sum(condition)等效于python numpy

时间:2019-05-01 12:24:41

标签: python matlab numpy

我正在尝试将一段matlab代码转换为python。

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

实际上返回每列中有条件的元素数。 numpy(python)中有任何等效的函数可以做到这一点吗?

2 个答案:

答案 0 :(得分: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

答案 1 :(得分:0)

尽管不是为此目的而设计的,但是可以选择替代方法

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)