Python numpy 非零 cumsum

Question

我想用numpy数组进行非零cumsum。只需跳过数组中的零并应用 cumsum 即可。假设我有一个 np.数组

a = np.array([1,2,1,2,5,0,9,6,0,2,3,0])

我的结果应该是

[1,3,4,6,11,0,20,26,0,28,31,0]

我已经尝试过了

a = np.cumsum(a[a!=0])

但结果是

[1,3,4,6,11,20,26,28,31]

有什么想法吗？

Answer 1

您需要屏蔽原始数组，以便仅覆盖非零元素:

In [9]:
a = np.array([1,2,1,2,5,0,9,6,0,2,3,0])
a[a!=0] = np.cumsum(a[a!=0])
a

Out[9]:
array([ 1,  3,  4,  6, 11,  0, 20, 26,  0, 28, 31,  0])

另一种方法是使用np.where:

In [93]:
a = np.array([1,2,1,2,5,0,9,6,0,2,3,0])
a = np.where(a!=0,np.cumsum(a),a)
a

Out[93]:
array([ 1,  3,  4,  6, 11,  0, 20, 26,  0, 28, 31,  0])

时间

In [91]:
%%timeit
a = np.array([1,2,1,2,5,0,9,6,0,2,3,0])
a[a!=0] = np.cumsum(a[a!=0])
a

The slowest run took 4.93 times longer than the fastest. This could mean that an intermediate result is being cached 
100000 loops, best of 3: 12.6 µs per loop

In [94]:    
%%timeit
a = np.array([1,2,1,2,5,0,9,6,0,2,3,0])
a = np.where(a!=0,np.cumsum(a),a)
a

The slowest run took 6.00 times longer than the fastest. This could mean that an intermediate result is being cached 
100000 loops, best of 3: 10.5 µs per loop

上面显示 np.where 比第一种方法稍快