Check equality of multiple elements in array

Question

I'm new to Python from Matlab.

I want to create a new variable from a subset of an existing numpy array based on equality to some condition specified by a third numpy array, an ID in this case.

This works fine for one equality.

new_x = old_x[someID == 1]

But if I try to extend it several equalities at once it no longer works:

new_x = old_x[someID == 1:3]

Ideally I want to be able to choose many equalities, like:

new_x = old_x[someID == 1:3,7]

I could loop through each number I want to check but is there a simpler way of doing this?

What is old_x? A numpy array?

Dani Mesejo
– Dani Mesejo

2021-10-07 11:40:04 +00:00
Commented Oct 7, 2021 at 11:40 — Dani Mesejo
– Dani Mesejo, Commented Oct 7, 2021 at 11:40
Yes, I've edited that in thanks.

N Blake
– N Blake

2021-10-07 11:42:06 +00:00
Commented Oct 7, 2021 at 11:42 — N Blake
– N Blake, Commented Oct 7, 2021 at 11:42

Dani Mesejo · Accepted Answer · 2021-10-07 11:50:35Z

5

You could use np.isin + np.r_:

import numpy as np

# for reproducible results
np.random.seed(42)

# toy data
old_x = np.random.randint(10, size=100)

# create new array by filtering on boolean mask
new_x = old_x[np.isin(old_x, np.r_[1:3,7])]

print(new_x)

Output

[7 2 7 7 7 2 1 7 1 2 2 2 1 1 1 7 2 1 7 1 1 1 7 7 1 7 7 7 7 2 7 2 2 7]

You could substitute np.r_ by something like [1, 2, 7] and use it as below:

new_x = old_x[np.isin(old_x, [1, 2, 7])]

Additionally if the array is 1-dimensional you could use np.in1d:

new_x = old_x[np.in1d(old_x, [1, 2, 7])]
print(new_x)

Output (from in1d)

[7 2 7 7 7 2 1 7 1 2 2 2 1 1 1 7 2 1 7 1 1 1 7 7 1 7 7 7 7 2 7 2 2 7]

edited Oct 7, 2021 at 11:50

answered Oct 7, 2021 at 11:44

Dani Mesejo

62.2k6 gold badges56 silver badges86 bronze badges

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

2 Comments

accdias Over a year ago

Just retracted my close vote because of your answer. Good one! It made me better understand the question.

N Blake Over a year ago

Using the np.r_ works perfectly thanks, seems like that might be useful little function. And thanks for generating toy data, I'll be sure to do that myself in fututre.

Collectives™ on Stack Overflow

Check equality of multiple elements in array

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related