How to identify elements that belong to a set in a numpy array in python

Question

Say I have a numpy array

A = numpy.array([-1, 1, 2, -2, 3, -3])

I would like to get all numbers whose squares equal 1 or 9 (so the expected result is [1, -1, 3, -3]). I tried A[A**2 in [1, 9]] but got an error. Is there any built-in function to handle this simple task without doing loops? Thanks.

Answer 1

numpy has a fucnction that does what you what called in1d:

import numpy

A = numpy.array([-1, 1, 2, -2, 3, -3])
mask = numpy.in1d(A**2, [1, 9])
print(mask)
# [ True  True False False  True  True]
print(A[mask])
# [-1  1  3 -3]

Answer 2

You can use numpy.logical_or :

>>> import numpy as np
>>> A[np.logical_or(A**2==1,A**2==9)]
array([-1,  1,  3, -3])