One way would be to sum with indexed on original array and then subtract out the summations caused by the last indexed ones by -1s
-
out = M[E].sum(1) - M[-1]*(E==-1).sum(1)[:,None]
Another way would be pad zeros at the end of M
, so that those -1
would index into those zeros and hence have no effect on the final sum after indexing -
M1 = np.vstack((M, np.zeros((1,M.shape[1]), dtype=M.dtype)))out = M1[E].sum(1)
If there is exactly one or lesser -1
per row in E
, we can optimize further -
out = M[E].sum(1)m = (E==-1).any(1)out[m] -= M[-1]
Another based on tensor-multiplication -
np.einsum('ij,kli->kj',M, (E[...,None]==np.arange(M.shape[1])))