如何在python中識別屬於numpy數組中的集合的元素

Question

說我有一個numpy數組

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

我想得到方塊等於1或9的所有數字（因此預期結果為[1, -1, 3, -3] ）。 我試過A[A**2 in [1, 9]]但是出錯了。 是否有任何內置函數來處理這個簡單的任務而不進行循環？ 謝謝。

Answer 1

numpy有一個功能，可以做你所謂的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

你可以使用numpy.logical_or ：

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

Answer 3

你可以使用filter ：

>>> filter(lambda x: x ** 2 in [1, 9], A)
[-1, 1, 3, -3]