Necesitamos usar una lista de índices que deben mantenerse ordenados en lugar de alimentar el parámetro k-ésimo como un escalar. Por lo tanto, para mantener la naturaleza ordenada en los primeros 5 elementos, en lugar de np.argpartition(a,5)[:5] , simplemente haga:

 np.argpartition(a,range(5))[:5]

Aquí hay una ejecución de muestra para aclarar las cosas:

 In [84]: a = np.random.rand(10) In [85]: a Out[85]: array([ 0.85017222, 0.19406266, 0.7879974 , 0.40444978, 0.46057793, 0.51428578, 0.03419694, 0.47708 , 0.73924536, 0.14437159]) In [86]: a[np.argpartition(a,5)[:5]] Out[86]: array([ 0.19406266, 0.14437159, 0.03419694, 0.40444978, 0.46057793]) In [87]: a[np.argpartition(a,range(5))[:5]] Out[87]: array([ 0.03419694, 0.14437159, 0.19406266, 0.40444978, 0.46057793])

Tenga en cuenta que argpartition tiene sentido en el aspecto de rendimiento, si buscamos obtener índices ordenados para un pequeño subconjunto de elementos, digamos k número de elementos, que es una pequeña fracción del número total de elementos.

Usemos un conjunto de datos más grande e intentemos obtener índices ordenados para todos los elementos para aclarar el punto mencionado anteriormente:

 In [51]: a = np.random.rand(10000)*100 In [52]: %timeit np.argpartition(a,range(a.size-1))[:5] 10 loops, best of 3: 105 ms per loop In [53]: %timeit a.argsort() 1000 loops, best of 3: 893 µs per loop

Por lo tanto, para ordenar todos los elementos, np.argpartition no es el camino a seguir.

Ahora, digamos que quiero obtener índices ordenados solo para los primeros 5 elementos con ese gran conjunto de datos y también mantener el orden para esos:

 In [68]: a = np.random.rand(10000)*100 In [69]: np.argpartition(a,range(5))[:5] Out[69]: array([1647, 942, 2167, 1371, 2571]) In [70]: a.argsort()[:5] Out[70]: array([1647, 942, 2167, 1371, 2571]) In [71]: %timeit np.argpartition(a,range(5))[:5] 10000 loops, best of 3: 112 µs per loop In [72]: %timeit a.argsort()[:5] 1000 loops, best of 3: 888 µs per loop

¡Muy útil aquí!

Dada la tarea de ordenar indirectamente un subconjunto (el top k , top significa primero en el orden de clasificación) hay dos soluciones integradas: argsort y argpartition cf. Respuesta de @Divakar.

Sin embargo, si el rendimiento es una consideración, entonces (dependiendo del tamaño de los datos y el subconjunto de interés) puede valer la pena resistir el "atractivo de la línea única", invertir una línea más y aplicar argsort en la salida de argpartition :

 >>> def top_k_sort(a, k): ... return np.argsort(a)[:k] ... >>> def top_k_argp(a, k): ... return np.argpartition(a, range(k))[:k] ... >>> def top_k_hybrid(a, k): ... b = np.argpartition(a, k)[:k] ... return b[np.argsort(a[b])] >>> k = 100 >>> timeit.timeit('f(a,k)', 'a=rng((100000,))', number = 1000, globals={'f': top_k_sort, 'rng': np.random.random, 'k': k}) 8.348663672804832 >>> timeit.timeit('f(a,k)', 'a=rng((100000,))', number = 1000, globals={'f': top_k_argp, 'rng': np.random.random, 'k': k}) 9.869098862167448 >>> timeit.timeit('f(a,k)', 'a=rng((100000,))', number = 1000, globals={'f': top_k_hybrid, 'rng': np.random.random, 'k': k}) 1.2305558240041137

argsort es O(n log n), argpartition con argumento de rango parece ser O(nk) (?), y argpartition + argsort es O(n + k log k)

Por lo tanto, en un régimen interesante n >> k >> 1, se espera que el método híbrido sea el más rápido

ACTUALIZACIÓN: versión ND:

 import numpy as np from timeit import timeit def top_k_sort(A,k,axis=-1): return A.argsort(axis=axis)[(*axis%A.ndim*(slice(None),),slice(k))] def top_k_partition(A,k,axis=-1): return A.argpartition(range(k),axis=axis)[(*axis%A.ndim*(slice(None),),slice(k))] def top_k_hybrid(A,k,axis=-1): B = A.argpartition(k,axis=axis)[(*axis%A.ndim*(slice(None),),slice(k))] return np.take_along_axis(B,np.take_along_axis(A,B,axis).argsort(axis),axis) A = np.random.random((100,10000)) k = 100 from timeit import timeit for f in globals().copy(): if f.startswith("top_"): print(f, timeit(f"{f}(A,k)",globals=globals(),number=10)*100)

Ejemplo de ejecución:

 top_k_sort 63.72379460372031 top_k_partition 99.30561298970133 top_k_hybrid 10.714635509066284

Describamos el método de partición de una manera simplificada que ayuda mucho a entender argpartition

Siguiendo el ejemplo de la imagen si ejecutamos C=numpy.argpartition(A, 3) C será el arreglo resultante de obtener la posición de cada elemento en B con respecto al arreglo A. es decir:

 Idx(z) = index of element z in array A then C would be C = [ Idx(B[0]), Idx(B[1]), Idx(B[2]), Idx(X), Idx(B[4]), ..... Idx(B[N]) ]

Como se mencionó anteriormente, este método es muy útil y es muy útil cuando tiene una matriz enorme y solo está interesado en un grupo seleccionado de elementos ordenados, no en la matriz completa.