Mask NumPy array and extract values where condition is met

Question

How can I mask two NumPy arrays properly? I want find pe values that are not equal to 255, for example. I also want my desired output array be the same size aspd and pe, i.e., (7, 7) and filled with 0's.
What is the most efficient way to achieve this?

import numpy as np

pd = np.random.randint(254,256, size=(7,7))
pe = np.random.randint(0,7, size=(7,7))

Desired output

[[6 6 0 0 6 6 1]
 [2 6 1 1 5 6 3]
 [3 4 6 6 3 5 6]
 [3 5 0 3 2 0 0]
 [0 3 6 1 3 6 1]
 [6 3 4 1 0 3 1]
 [6 0 4 2 2 6 4]]

Many thanks

So for each index in the output array, you'd like pe[index] if pd[index] != 255 else 0, right? — dsm
– dsm, Commented Jun 6, 2022 at 12:35
pe * (pd == 255), IIUC. Assuming you meant that are not equal to 255 in pd — Michael Szczesny
– Michael Szczesny, Commented Jun 6, 2022 at 12:36
Which is slightly faster than np.where(pd == 255, 0, pe) which I was going to propose. Kudos! :) — dsm
– dsm, Commented Jun 6, 2022 at 12:38
Note that we optimized significantly np.where in the very last version (1.23) of Numpy so np.where(pd == 255, 0, pe) may now be faster than pe * (pd == 255) or at least close to that. — Jérôme Richard
– Jérôme Richard, Commented Jun 6, 2022 at 13:42

AboAmmar · Accepted Answer · 2022-06-06 14:04:06Z

1

Logical indexing seems the simplest of all options.

import numpy as np

pd = np.random.randint(254,256, size=(7,7))
pe = np.random.randint(0,7, size=(7,7))

pe[pd == 255] = 0

[[3 6 0 2 0 0 0]
 [0 3 4 5 2 0 5]
 [0 0 6 0 1 0 5]
 [0 3 0 4 0 6 0]
 [2 0 0 0 0 0 0]
 [2 0 4 0 0 0 5]
 [0 0 3 0 2 4 0]]

Based on your data size, you may try other options:

pe = np.where(pd == 255, 0, pe)
# OR
pe = pe * (pd == 255)

but I guess indexing is still simple and fast.

answered Jun 6, 2022 at 14:04

AboAmmar

5,5892 gold badges15 silver badges25 bronze badges

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

1 Comment

Rimo Over a year ago

Many thanks for a clear reply, AboAmmar!

Collectives™ on Stack Overflow

Mask NumPy array and extract values where condition is met

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