CODEWITHSUNDEEP
pythonnumpy
Bubaya
Bubaya

Reputation: 823

Numpy: sums 1D array indexed by rows of 2D boolean array

Assume that I have an (m,)-array a and an (n, m)-array b of booleans. For each row b[i] of b, I want to take np.sum(a, where=b[i]), which should return an (n,)-array. I could do the following:

a = np.array([1,2,3])
b = np.array([[True, False, False], [True, True, False], [False, True, True]])
c = np.array([np.sum(a, where=r) for r in b])
# c is [1,3,5]

but this seems quite unelegant to me. I would have hoped that broadcasting magic makes something like

c = np.sum(a, where=b)
# c is 9

work, but apparently, np.sum then sums over the rows of b, which I do not want. Is there a numpy-inherent way of achieving the desired behavour with np.sum (or any ufunc.reduce)?

Upvotes: 0

Views: 56

Answers (1)

yann ziselman
yann ziselman

Reputation: 2002

How about:

a = np.array([1,2,3])
b = np.array([[True, False, False], [True, True, False], [False, True, True]])
c = np.sum(a*b, axis = 1)

output:

array([1, 3, 5])

Upvotes: 2

Related Questions