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

Question

我有一个包含10万行的2-D numpy数组。我需要返回其中的一部分行（我需要执行很多次这样的操作，因此效率很重要）。

一个模拟例子如下：

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

And I want to return:

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

- mishaF

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

- JoshAdel

---

没错，这很酷，但在我的情况下，我需要匹配值。例如，b就像标识符，而不是索引。我会编辑问题以澄清这一点。 - mishaF

即使涉及广播，(a - b) == 0与a == b是相同的。 - Sven Marnach

@JoshAdel 非常感谢！幸运的是，我的 b 数组是整数，所以在浮点数问题上应该没问题。 - mishaF

@Josh: 我们两个答案都让我不满的是它们的复杂度都是 O(len(a)*len(b))，理论上只需要 O((len(a)+len(b))*log(len(b))) 就足够了（对 b 进行排序，并对 a[:,0] 的每个元素进行二进制搜索）。有什么想法可以改进吗？我们能使用 searchsorted() 吗？ - Sven Marnach

@Sven：很好的建议 - np.searchsorted 在这种情况下易于应用，而且速度显着更快。 - JoshAdel

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

- Sven Marnach

---

致敬Sven：我认为他的方法比我的解决方案快大约1.6倍。 - JoshAdel

@Josh：感谢你计时！你已经得到我的+1了，因为你首先提供了一个可行的答案。 :) - Sven Marnach

正如我在帖子的编辑3中所示，您可以直接使用searchsorted。值得注意的是，您提供的两种解决方案都只提取了b中唯一的条目，因此如果这对于OP很重要，那么这也是需要考虑的。 - JoshAdel