how to remove some of one array's elements according to other array's value in numpy

Question

In numpy,

a = np.array([[0, 0, 0, 1, 1, 1, 0], [0, 0, 1, 1, 1, 0, 0]])
b = np.array([[5, 2, 0, 1, 0, 2, 3], [0, 0, 0, 1, 5, 3, 0]])

I want to remove some of b‘s elements according to a‘s values. I want to keep the elements of b only when the value of a‘s matching elements is 1, and discard the other elements of b.
In this example what I want is:

>>> [[1, 0, 2], [0, 1, 5]]

How can I get this result?

Asked By: June Yoon

||

Source

Answer 1

You can apply the next line:

result = b[a == 1].reshape((2, 3))

Update
In case of dynamic amount of ones you can use the next code:

mask = a == 1
result = [b_row[mask_row] for b_row, mask_row in zip(b, mask)]

or

result = []
for i in range(len(a)):
    mask = a[i] == 1
    result.append(b[i][mask].tolist())

Answered By: Aksen P

Answer 2

You can use a boolean mask to index into b and finally reshape the result.

res = b[a == 1].reshape(a.shape[0], -1)
print(res)

Or

res = [list(np.array(b_row)[a_row == 1]) for a_row, b_row in zip(a, b)]
print(res)

[[1 0 2]
 [0 1 5]]

Answered By: Jamiu S.

Answer 3

considering What if the number of elements with a value of 1 is different each time? It's not always same 3.

maybe that’s sufficient solution:

result = [
    [
        element_b
        for element_a, element_b in zip(list_a, list_b)
        if element_a == 1
    ]
    for list_a, list_b in zip(a,b)
]

Answered By: Jaroszewski Piotr

how to remove some of one array's elements according to other array's value in numpy

Question:

Answers: