Ordena una matriz numpy por otra

Tengo una matriz que determina un orden de elementos:

order = [3, 1, 4, 2] 

Y luego quiero ordenar otra matriz más grande (que contiene solo esos elementos):

 a = np.array([4, 2, 1, 1, 4, 3, 1, 3]) 

de tal manera que los elementos que vienen primero en order están primero en los resultados, etc.
En Python directo, haría esto con una función clave:

 sorted(a, key=order.index) [3, 3, 1, 1, 1, 4, 4, 2] 

¿Cómo puedo hacer esto (eficientemente) con numpy? ¿Existe una noción similar de “función clave” para matrices numpy?

Caso específico: Ints

Para ints , podríamos usar bincount

 np.repeat(order,np.bincount(a)[order]) 

Ejecución de la muestra

 In [146]: sorted(a, key=order.index) Out[146]: [3, 3, 1, 1, 1, 4, 4, 2] In [147]: np.repeat(order,np.bincount(a)[order]) Out[147]: array([3, 3, 1, 1, 1, 4, 4, 2]) 

Caso genérico

Enfoque # 1

Generalizando para todos los tipos con bincount

 # https://stackoverflow.com/a/41242285/ @Andras Deak def argsort_unique(idx): n = idx.size sidx = np.empty(n,dtype=int) sidx[idx] = np.arange(n) return sidx sidx = np.argsort(order) c = np.bincount(np.searchsorted(order,a,sorter=sidx)) out = np.repeat(order, c[argsort_unique(sidx)]) 

Enfoque # 2-A

Con np.unique y searchsorted para el caso cuando todos los elementos de order están en a

 unq, count = np.unique(a, return_counts=True) out = np.repeat(order, count[np.searchsorted(unq, order)]) 

Enfoque # 2-B

Para cubrir todos los casos, necesitamos un paso adicional:

 unq, count = np.unique(a, return_counts=1) sidx = np.searchsorted(unq, order) out = np.repeat(order, np.where(unq[sidx] == order,count[sidx],0)) 

Sobre la base de la solución de @Divakar, puede contar cuántas veces ocurre cada elemento y luego repetir los elementos ordenados que muchas veces:

 c = Counter(a) np.repeat(order, [c[v] for v in order]) 

(Si lo desea, puede vectorizar la búsqueda de conteo). Me gusta esto porque es el tiempo lineal, incluso si no es puro número.

Supongo que un equivalente puro numpy se vería así:

 count = np.unique(a, return_counts=True)[1] np.repeat(order, count[np.argsort(np.argsort(order))]) 

Pero eso es menos directo, más código y demasiados tipos. 🙂

Esta es una conversión bastante directa de su enfoque de Python puro en números. La idea clave es reemplazar la función order.index con una búsqueda en un vector ordenado. No estoy seguro si esto es más simple o más rápido que la solución que se te ocurrió, pero puede generalizarse a otros casos.

 import numpy as np order = np.array([3, 1, 4, 2]) a = np.array([4, 2, 1, 1, 4, 3, 1, 3]) # create sorted lookup vectors ord = np.argsort(order) order_sorted = order[ord] indices_sorted = np.arange(len(order))[ord] # lookup the index in `order` for each value in the `a` vector a_indices = np.interp(a, order_sorted, indices_sorted).astype(int) # sort `a` using the retrieved index values a_sorted = a[np.argsort(a_indices)] a_sorted # array([3, 3, 1, 1, 1, 4, 4, 2]) 

Esta es una forma más directa (basada en esta pregunta ), pero parece ser aproximadamente 4 veces más lenta que el enfoque np.interp :

 lookup_dict = dict(zip(order, range(len(order)))) indices = np.vectorize(lookup_dict.__getitem__)(a) a_sorted = a[np.argsort(indices)]