Scipy Sparse Cumsum
How about doing this instead
a = np.array([[0,0,1,2,0,3,0,4], [1,0,0,2,0,3,4,0]], dtype=int)b = a.copy()b[b > 0] = 1z = np.cumsum(a,axis=1)print(z*b)Yields
array([[ 0, 0, 1, 3, 0, 6, 0, 10], [ 1, 0, 0, 3, 0, 6, 10, 0]])Doing sparse
def sparse(a): a = scipy.sparse.csr_matrix(a) indptr = a.indptr data = a.data for i in range(a.shape[0]): st = indptr[i] en = indptr[i + 1] np.cumsum(data[st:en], out=data[st:en])In[1]: %timeit sparse(a)10000 loops, best of 3: 167 µs per loopUsing multiplication
def mult(a): b = a.copy() b[b > 0] = 1 z = np.cumsum(a, axis=1) z * bIn[2]: %timeit mult(a)100000 loops, best of 3: 5.93 µs per loop