numpy.searchsorted与2D数组

Question

我有这个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中更有效率地做到这一点？

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