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

Question

我有两个nparrays

A 维度 (60000,784)

和

B 维度 (60000, 1)

它们都存储整数。

我想创建一个包含 A 行的 nparray，这样 B 的相应行要么是某个数字 X 要么是某个数字 Y。我还想要返回包含 B 的相应行的第二个 nparray。

例子:

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 创建 bool 索引。

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]