python - 从二维数组中提取指定行的最有效方法？

Question

我有一个包含 100,000 多行的二维 numpy 数组。我需要返回这些行的一个子集(我需要执行该操作 1,000 次，因此效率很重要)。

一个模拟的例子是这样的:

import numpy as np
a = np.array([[1,5.5],
             [2,4.5],
             [3,9.0],
             [4,8.01]])
b = np.array([2,4])

所以...我想从 a 返回数组，其中第一列中的行由 b 标识:

c=[[2,4.5],
   [4,8.01]]

当然，不同之处在于 a 和 b 中都有更多的行，所以我想避免循环。另外，我尝试制作字典并使用 np.nonzero，但仍然有点难过。

提前感谢您的任何想法!

编辑:请注意，在这种情况下，b 是标识符而不是索引。这是一个修改后的例子:

import numpy as np
a = np.array([[102,5.5],
             [204,4.5],
             [343,9.0],
             [40,8.01]])
b = np.array([102,343])

我想返回:

c = [[102,5.5],
     [343,9.0]]

Answer 1

编辑: 删除了我原来的答案，因为这是对问题的误解。而是尝试:

ii = np.where((a[:,0] - b.reshape(-1,1)) == 0)[1]
c = a[ii,:]

我正在做的是使用广播从 a 中减去 b 的每个元素，然后在该数组中搜索表示匹配的零。这应该可行，但您应该小心比较 float ，尤其是当 b 不是整数数组时。

EDIT 2 感谢 Sven 的建议，您可以试试这个稍微修改过的版本:

ii = np.where(a[:,0] == b.reshape(-1,1))[1]
c = a[ii,:]

它比我原来的实现要快一点。

编辑 3 目前最快的解决方案(比 Sven 的第二个大型阵列解决方案快 10 倍)是:

c = a[np.searchsorted(a[:,0],b),:]

假设 a[:,0] 已排序并且 b 的所有值都出现在 a[:,0] 中。