Векторизованный searchsorted numpy

Question

Предположим, что у меня есть два массива A и B, где A и B являются m x n. Моя цель теперь, для каждой строки A и B, найти, где я должен вставить элементы строки i из A в соответствующую строку B. То есть, я хочу применить np.digitize или np.searchsorted к каждой строке A и B.

Мое наивное решение состоит в том, чтобы просто перебирать строки. Однако это слишком медленно для моего приложения. Поэтому мой вопрос: есть ли векторизованная реализация любого алгоритма, который мне не удалось найти?

Answer 1

Мы можем добавить каждую смещение ряда по сравнению с предыдущей строкой. Мы использовали бы такое же смещение для обоих массивов. Идея состоит в использовании np.searchsorted на сплющенной версии входных массивов после этого, и поэтому каждая строка из b будет ограничена для поиска отсортированных позиций в соответствующей строке в a. Кроме того, чтобы заставить его работать и для отрицательных чисел, нам просто нужно компенсировать минимальные числа.

Таким образом, мы имели бы векторизованную реализацию как так -

def searchsorted2d(a,b): 
    m,n = a.shape 
    max_num = np.maximum(a.max() - a.min(), b.max() - b.min()) + 1 
    r = max_num*np.arange(a.shape[0])[:,None] 
    p = np.searchsorted((a+r).ravel(), (b+r).ravel()).reshape(m,-1) 
    return p - n*(np.arange(m)[:,None]) 

время выполнения теста -

In [173]: def searchsorted2d_loopy(a,b): 
    ...:  out = np.zeros(a.shape,dtype=int) 
    ...:  for i in range(len(a)): 
    ...:   out[i] = np.searchsorted(a[i],b[i]) 
    ...:  return out 
    ...: 

In [174]: # Setup input arrays 
    ...: a = np.random.randint(11,99,(10000,20)) 
    ...: b = np.random.randint(11,99,(10000,20)) 
    ...: a = np.sort(a,1) 
    ...: b = np.sort(b,1) 
    ...: 

In [175]: np.allclose(searchsorted2d(a,b),searchsorted2d_loopy(a,b)) 
Out[175]: True 

In [176]: %timeit searchsorted2d_loopy(a,b) 
10 loops, best of 3: 28.6 ms per loop 

In [177]: %timeit searchsorted2d(a,b) 
100 loops, best of 3: 13.7 ms per loop