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

2024-03-29 01:25:13 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个有100000多行的二维numpy数组。我需要返回这些行的一个子集(我需要执行这些操作1000次,所以效率很重要)。在

模型示例如下:

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标识行:

^{pr2}$

当然,区别在于a和b中有更多的行,所以我希望避免循环。而且,我还玩着编字典和np.非零但我还是有点困惑。在

提前感谢您的任何想法!在

编辑:注意,在本例中,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]]

Tags: 模型importnumpy编辑示例字典asnp
2条回答

一种更简洁的方法是

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

浮点比较的常见注意事项适用。在

编辑:如果b不太小,则以下略显古怪的解决方案执行得更好:

^{pr2}$

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

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

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

编辑2多亏了斯文的建议,你可以试试这个稍作修改的版本:

^{pr2}$

它比我最初的实现要快一点。在

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

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

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

相关问题 更多 >