Question

我有一个拥有100,000多行的2-D 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])

所以...我想从第一列中用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,:]

我正在做的是使用广播从b中减去a的每个元素，然后在该数组中搜索表示匹配的零。这应该有效，但是你应该对浮点数进行比较时要小心，特别是如果b不是一个整数数组。

编辑2 感谢Sven的建议，您可以尝试这个稍微修改过的版本：

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

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

EDIT 3 迄今为止最快的解决方案（比Sven的第二个大型阵列解决方案快10倍）是：

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

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

Answer 2

更简洁的方法是

c = a[(a[:,0] == b[:,None]).any(0)]

浮点比较的常见警告适用。

修改：如果b不是太小，以下略显古怪的解决方案效果会更好：

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