从二维数组中提取指定行的最有效方法是什么？

Question

我有一个2-D numpy阵列,有100,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,:]

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

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

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

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