Как искать отсортированные 2D-массивы для индексов вставки строк в Python?

Question

Я ищу в отсортированном массиве правильные индексы вставки новых данных, чтобы они оставались отсортированными.Хотя searchsorted2d от @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).