numpy.searchsorted 与二维数组

Posted 2023-02-23

【中文标题】numpy.searchsorted 与二维数组

【英文标题】：numpy.searchsorted with 2D array

【发布时间】：2014-04-29 04:33:45

【问题描述】：

我有这个 numpy 数组，其中每一行中的值将始终被排序并单调递增：

a = np.array([[1, 2, 3, 4, 8],
       [2, 5, 6, 7, 8],
       [5, 7, 11, 12, 13]])

我想为每一行搜索以下值（未排序或单调）：

b = np.array([4.5, 2.3, 11.6])

以便我得到以下答案：

[4, 1, 3]

但是，searchsorted 不支持这个（感觉需要一个axis 关键字）。

对于一个非常大的数组，有没有一种有效的方法可以做到这一点？显然，使用for 循环我可以像这样索引数组a 和b：

for i in np.arange(np.alen(a)):
     print a[i].searchsorted(b[i])

但是当a 很大时，这会很慢。

有没有办法在 numpy 中更有效地做到这一点？

【问题讨论】：

这道题为什么是负分？

【参考方案1】：

您可以在 ravel/flattened 数组上搜索排序：

In [11]: np.searchsorted(a.ravel(), b)
Out[11]: array([3, 6])

然后您可以对结果使用 divmod（获取行和列）：

In [12]: divmod(np.searchsorted(a.ravel(), b), a.shape[1])
Out[12]: (array([0, 1]), array([3, 1]))

【讨论】：

这要求每一行的值都大于任何前一行，这是 OP 示例的情况，但对于一般数组来说似乎是一个强条件。

@Andy - 问题中示例的优雅解决方案，但正如 Jamie 所说，它不适用于 b 不是单调上升的情况。我已经更新了这个问题，所以它更清楚了。谢谢。

@KernowBunney 啊，好吧，这样效率更高但类似于np.diag(np.apply_along_axis(np.searchsorted, 1, a, b))。使用 cython 是提高 for 循环性能的一个简单选择。

@Andy Hayden - 谢谢 - apply_along_axis 有效，并且至少比 Python 循环更简洁/类似 Numpy。我将调查 Cython。干杯！

@AndyHayden，apply_along_axis 很棒！