argsort para un ndarray multidimensional

Estoy tratando de obtener los índices para ordenar una matriz multidimensional por el último eje, por ejemplo

>>> a = np.array([[3,1,2],[8,9,2]]) 

Y me gustaria indices de tal manera,

 >>> a[i] array([[1, 2, 3], [2, 8, 9]]) 

Basándome en la documentación de numpy.argsort , pensé que debería hacer esto, pero recibo el error:

 >>> a[np.argsort(a)] IndexError: index 2 is out of bounds for axis 0 with size 2 

Edición: necesito reorganizar otras matrices de la misma forma (por ejemplo, una matriz b tal que a.shape == b.shape ) de la misma manera … para que

 >>> b = np.array([[0,5,4],[3,9,1]]) >>> b[i] array([[5,4,0], [9,3,1]]) 

Solución:

 >>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)] array([[1, 2, 3], [2, 8, 9]]) 

Lo entendiste bien, aunque no lo describiría como una trampa en la indexación.

Tal vez esto ayude a aclararlo:

 In [544]: i=np.argsort(a,axis=1) In [545]: i Out[545]: array([[1, 2, 0], [2, 0, 1]]) 

i es el orden que queremos, para cada fila. Es decir:

 In [546]: a[0, i[0,:]] Out[546]: array([1, 2, 3]) In [547]: a[1, i[1,:]] Out[547]: array([2, 8, 9]) 

Para realizar ambos pasos de indexación a la vez, tenemos que usar un índice de ‘columna’ para la primera dimensión.

 In [548]: a[[[0],[1]],i] Out[548]: array([[1, 2, 3], [2, 8, 9]]) 

Otra matriz que podría emparejarse con i es:

 In [560]: j=np.array([[0,0,0],[1,1,1]]) In [561]: j Out[561]: array([[0, 0, 0], [1, 1, 1]]) In [562]: a[j,i] Out[562]: array([[1, 2, 3], [2, 8, 9]]) 

Si i identifica la columna para cada elemento, j especifica la fila para cada elemento. La matriz de columnas [[0],[1]] funciona igual de bien porque se puede transmitir contra i .

pienso

 np.array([[0], [1]]) 

como ‘mano corta’ para j . Juntos, definen la fila y la columna de origen de cada elemento de la nueva matriz. Trabajan juntos, no secuencialmente.

La asignación completa de a a la nueva matriz es:

 [a[0,1] a[0,2] a[0,0] a[1,2] a[1,0] a[1,1]] 

 def foo(a): i = np.argsort(a, axis=1) return (np.arange(a.shape[0])[:,None], i) In [61]: foo(a) Out[61]: (array([[0], [1]]), array([[1, 2, 0], [2, 0, 1]], dtype=int32)) In [62]: a[foo(a)] Out[62]: array([[1, 2, 3], [2, 8, 9]]) 

Encontré la respuesta aquí , con alguien teniendo el mismo problema. La clave es simplemente engañar a la indexación para que funcione correctamente …

 >>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)] array([[1, 2, 3], [2, 8, 9]]) 

También puede utilizar linear indexing , que podría ser mejor con el rendimiento, por lo que –

 M,N = a.shape out = b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)] 

Entonces, a.argsort(1)+(np.arange(M)[:,None]*N) son básicamente los índices lineales que se utilizan para asignar b para obtener la salida ordenada deseada para b . Los mismos índices lineales también podrían usarse en a para obtener la salida ordenada para a .

Ejecución de la muestra

 In [23]: a = np.array([[3,1,2],[8,9,2]]) In [24]: b = np.array([[0,5,4],[3,9,1]]) In [25]: M,N = a.shape In [26]: b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)] Out[26]: array([[5, 4, 0], [1, 3, 9]]) 

Pruebas de tiempo

 In [27]: a = np.random.rand(1000,1000) In [28]: b = np.random.rand(1000,1000) In [29]: M,N = a.shape In [30]: %timeit b[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)] 10 loops, best of 3: 133 ms per loop In [31]: %timeit b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)] 10 loops, best of 3: 96.7 ms per loop 

Las respuestas anteriores están ahora un poco desactualizadas, ya que se agregó nueva funcionalidad en numpy 1.15 para que sea más simple; take_along_axis ( https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.take_along_axis.html ) le permite hacer:

 >>> a = np.array([[3,1,2],[8,9,2]]) >>> np.take_along_axis(a, a.argsort(axis=-1), axis=-1) array([[1 2 3] [2 8 9]])