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

【问题标题】：Most efficient way to pull specified rows from a 2-d array?从二维数组中提取指定行的最有效方法？
【发布时间】：2011-03-31 19:41:29
【问题描述】：

我有一个包含 100,000 多行的二维 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 标识：

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]]

【问题讨论】：

【解决方案1】：

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

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

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

EDIT 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]中。

【讨论】：

对 - 这很酷，但就我而言，我需要匹配这些值。例如， b 就像标识符，而不是索引。我将编辑问题以澄清这一点。
(a - b) == 0 与 a == b 相同，即使涉及广播。
@JoshAdel 非常感谢！幸运的是，我的 b 数组是整数，所以我在浮点问题上应该没问题。
@Josh：让我对我们的两个答案感到不满的是复杂性是O(len(a)*len(b))，理论上O((len(a)+len(b))*log(len(b))) 就足够了（对b 进行排序并对@ 的每个元素进行二分搜索987654334@)。任何想法如何改善这一点？我们可以使用searchsorted()吗？
@Sven: 好电话 - np.searchsorted 很容易应用于这种情况，而且速度明显更快

【解决方案2】：

一种更简洁的方法是

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

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

编辑：如果b不是太小，下面这个稍微古怪的解决方案表现更好：

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

【讨论】：

对 Sven 的支持：我认为他的方法比我的解决方案快 1.6 倍。
@Josh：感谢您安排时间！无论如何，您首先提供了一个有效的答案，因此您得到了我的 +1。 :)
如我帖子编辑3所示，可以直接使用searchsorted。还值得注意的是，您的两个解决方案都只提取 b 中的唯一条目，所以如果这对 OP 很重要，那么这也是一个考虑因素。