如何使用 numpy 方法根据另一个 np 数组的条件对一个 np 数组的某些行执行操作？

Question

For example, I have one np array A = [[30, 60, 50...], [ 15, 20, 18...], [21, 81, 50...]...] of size (N, 10) .

And I have another np array B = [1, 1, 0...] of size (N, ) .

I want to do operations Eg I want all the sums of each column in A but only for rows where B==1. How would I do that without using any loops and just numpy methods?

So if I want sum of columns in A for indices where B == 1:

result = 30 + 15 because the first two indices in B are 1 but the third index is 0 so I wouldn't include it in the sum.

Answer 1

Use np.compress and sum along axis=0

>>> A = [[30, 60, 50], [ 15, 20, 18], [21, 81, 50]]
>>> B = [1, 1, 0]
>>> np.compress(B, A, axis=0).sum(0)
array([45, 80, 68])

If array, use np.nonzero on B :

>>> A = np.array([[30, 60, 50], [ 15, 20, 18], [21, 81, 50]])
>>> A[np.nonzero(B)].sum(0)
array([45, 80, 68])

Another way:

>>> A[B.astype(bool)].sum(0)
array([45, 80, 68])

If you want 0 s:

>>> np.compress(B==0, A, axis=0).sum(0)
# Or,
>>> A[np.nonzero(B==0)].sum(0)
# Or,
>>> A[~B.astype(bool)].sum(0)

If you want both 1 s and 0 s, obviously:

>>> A.sum(0)

Answer 2

You can convert B to bool type and mask A . Then you can get the sum along columns.

A = np.array([[30, 60, 50], [ 15, 20, 18], [21, 81, 50]])
B = np.array([1, 1, 0])

A[B.astype(np.bool)].sum(axis=0)

array([45, 80, 68])