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

Question

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)?

yann ziselman · Accepted Answer · 2021-06-24 12:03:44Z

2

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])

answered Jun 24, 2021 at 12:03

yann ziselman

2,0027 silver badges21 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Bubaya Over a year ago

OK, that was too easy.

Collectives™ on Stack Overflow

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

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related