如何在Python中搜索排序的2D数组中的行插入索引？

Question

我在一个排序数组中搜索新数据的正确插入索引，以便它保持排序。尽管@Divakar的^{}可以很好地处理列插入，但它不能处理行。有没有一种方法可以沿着行执行相同的操作？你知道吗

想到的第一个想法是调整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）。你知道吗