python - numpy.searchsorted 二维数组

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