在列约束下选择nparray的某些行

Question

我有两个nparrays

A尺寸（60000,784）

和

B尺寸（60000,1）

他们都存储整数。

我想创建一个包含来自A的行的nparray，使得B的相应行是某个数字X或某个数字Y.我还想返回第二个nparray包含B的相应行。

示例：

A

1 2
3 4
5 6
8 9


B

2
5
6
1

假设我想要A的行，使得B的相应行为1或5.然后返回的第一个nparray将是以下内容：

3 4
8 9

，第二个nparray将如下：

5
1

使用nparrays可以

这样吗？

提前谢谢

Answer 1

使用np.in1d创建一个布尔索引。

A = np.array([[1, 2], [3, 4], [5, 6], [8, 9]])
B = np.array([2,5,6,1])

index = np.in1d(B, [1,5])

print A[index]

print B[index]