python - Búsqueda vectorizada clasificada numpy

performance vectorization

Question

Suponga que tengo dos matrices A y B , donde A y B son mxn . Mi objetivo ahora es, para cada fila de A y B , encontrar dónde debo insertar los elementos de la fila i de A en la fila correspondiente de B Es decir, deseo aplicar np.digitize o np.searchsorted a cada fila de A y B

Mi solución ingenua es simplemente iterar sobre las filas. Sin embargo, esto es demasiado lento para mi aplicación. Por lo tanto, mi pregunta es: ¿hay una implementación vectorizada de cualquiera de los algoritmos que no he logrado encontrar?

Answer 1

La solución proporcionada por @Divakar es ideal para datos enteros, pero tenga cuidado con los problemas de precisión para los valores de coma flotante, especialmente si abarcan varios órdenes de magnitud (por ejemplo, [[1.0, 2,0, 3.0, 1.0e+20],...] ). En algunos casos, r puede ser tan grande que al aplicar a+r y b+r borra los valores originales en los que está tratando de ejecutar la searchsorted , y solo está comparando r con r .

Para hacer que el enfoque sea más robusto para los datos de punto flotante, puede incrustar la información de la fila en las matrices como parte de los valores (como un tipo estructurado) y ejecutar la búsqueda ordenada en estos tipos estructurados.

def searchsorted_2d (a, v, side=''left'', sorter=None):
  import numpy as np

  # Make sure a and v are numpy arrays.
  a = np.asarray(a)
  v = np.asarray(v)

  # Augment a with row id
  ai = np.empty(a.shape,dtype=[(''row'',int),(''value'',a.dtype)])
  ai[''row''] = np.arange(a.shape[0]).reshape(-1,1)
  ai[''value''] = a

  # Augment v with row id
  vi = np.empty(v.shape,dtype=[(''row'',int),(''value'',v.dtype)])
  vi[''row''] = np.arange(v.shape[0]).reshape(-1,1)
  vi[''value''] = v

  # Perform searchsorted on augmented array.
  # The row information is embedded in the values, so only the equivalent rows 
  # between a and v are considered.
  result = np.searchsorted(ai.flatten(),vi.flatten(), side=side, sorter=sorter)

  # Restore the original shape, decode the searchsorted indices so they apply to the original data.
  result = result.reshape(vi.shape) - vi[''row'']*a.shape[1]

  return result

Editar: ¡ El momento en este enfoque es abismal!

In [21]: %timeit searchsorted_2d(a,b)
10 loops, best of 3: 92.5 ms per loop

Sería mejor simplemente usando el map sobre la matriz:

In [22]: %timeit np.array(list(map(np.searchsorted,a,b)))
100 loops, best of 3: 13.8 ms per loop

Para datos enteros, el enfoque de @ Divakar sigue siendo el más rápido:

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

Answer 2

Podemos agregar a cada fila algún desplazamiento en comparación con la fila anterior. Usaríamos el mismo desplazamiento para ambas matrices. La idea es usar np.searchsorted en una versión aplanada de matrices de entrada a partir de entonces y, por lo tanto, cada fila de b se restringiría para encontrar posiciones ordenadas en la fila correspondiente en a . Además, para que funcione también con números negativos, solo necesitamos compensar los números mínimos también.

Entonces, tendríamos una implementación vectorizada así:

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

Prueba de tiempo de ejecución

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