如何在Python中搜索2D数组以查找行的插入索引？

Question

我正在搜索已排序的数组，以查找新数据的正确插入索引，以使其保持已排序状态。尽管@Divakar的searchsorted2d在插入列时效果很好，但不能在行中使用。有没有一种方法可以执行相同的操作，却沿行呢？

想到的第一个想法是使searchsorted2d适应所需的行为。但是，这似乎并不像看起来那样容易。这是我尝试进行调整的尝试，但是当axis设置为0时仍然无法使用。

import numpy as np

# By Divakar
# See https://stackoverflow.com/a/40588862
def searchsorted2d(a, b, axis=0):
    shape = list(a.shape)
    shape[axis] = 1
    max_num = np.maximum(a.max() - a.min(), b.max() - b.min()) + 1
    r = np.ceil(max_num) * np.arange(a.shape[1-axis]).reshape(shape)
    p = np.searchsorted((a + r).ravel(), (b + r).ravel()).reshape(b.shape)
    return p #- a.shape[axis] * np.arange(a.shape[1-axis]).reshape(shape)

axis = 0 # Operate along which axis?
n = 16   # vector size

# Initial array
a = np.random.rand(n).reshape((n, 1) if axis else (1, n))

insert_into_a = np.random.rand(n).reshape((n, 1) if axis else (1, n))
indices = searchsorted2d(a, insert_into_a, axis=axis)
a = np.insert(a, indices.ravel(), insert_into_a.ravel()).reshape(
              (n, -1) if axis else (-1, n))

assert(np.all(a == np.sort(a, axis=axis))), 'Failed :('
print('Success :)')

我希望断言在两种情况下都能通过（axis = 0和axis = 1）。