python - 检查数组中多个元素的相等性

Question

我是从 Matlab 开始接触 Python 的新手。

我想根据与第三个 numpy 数组(在本例中为 ID)指定的某些条件的相等性，从现有 numpy 数组的子集创建一个新变量。

这对于一个等式来说效果很好。

new_x = old_x[someID == 1]

但是如果我尝试一次扩展它几个等式，它就不再起作用:

new_x = old_x[someID == 1:3]

理想情况下，我希望能够选择许多等式，例如:

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

我可以循环遍历我想要检查的每个数字，但是有更简单的方法吗？

Answer 1

您可以使用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)

输出

[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]

您可以将 np.r_ 替换为 [1, 2, 7] 并按如下方式使用:

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

此外，如果数组是一维的，您可以使用 np.in1d :

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

输出 (来自 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]