Не могу понять вывод numpart argty

Я пытаюсь использовать arpgpartition из numpy, но кажется, что что-то пошло не так, и я не могу понять это. Вот что происходит:

Это первые 5 элементов отсортированного массива norms

np.sort(norms)[:5]
array([ 53.64759445,  54.91434479,  60.11617279,  64.09630585,  64.75318909], dtype=float32)

Но когда я использую indices_sorted = np.argpartition(norms, 5)[:5]

norms[indices_sorted]
array([ 60.11617279,  64.09630585,  53.64759445,  54.91434479,  64.75318909], dtype=float32)

Когда я думаю, что должен получить тот же результат, что и отсортированный массив?

Он отлично работает, когда я использую 3 как параметр indices_sorted = np.argpartition(norms, 3)[:3]

norms[indices_sorted]
array([ 53.64759445,  54.91434479,  60.11617279], dtype=float32)

Это не имеет большого смысла для меня, надеясь, что кто-то может предложить некоторое понимание?

РЕДАКТИРОВАТЬ: перефразируйте этот вопрос, как имеет смысл аргумент, сохраняющий порядок k секционированных элементов.

Ответ 1

Нам нужно использовать список индексов, которые должны храниться в отсортированном порядке, а не кормить kth param в качестве скаляра. Таким образом, чтобы сохранить отсортированный характер по первым элементам 5 вместо np.argpartition(a,5)[:5], просто сделайте -

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

Вот пример, чтобы сделать все понятным -

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

Обратите внимание, что argpartition имеет смысл в аспекте производительности, если мы хотим получить отсортированные индексы для небольшого подмножества элементов, скажем, k количество элем, которое составляет небольшую часть от общего числа элемен.

Позвольте использовать больший набор данных и попытайтесь получить отсортированные индексы для всех элемен- тов, чтобы сделать вышеописанную точку ясной -

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

Таким образом, чтобы отсортировать все элементы, np.argpartition не подходит.

Теперь, скажем, я хочу получить отсортированные индексы только для первых 5 элементов с этим большим набором данных, а также сохранить порядок для них -

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

Очень полезно здесь!

Ответ 2

Учитывая задачу неправильной сортировки подмножества (верхний k, верхний означает первый в порядке сортировки), есть два встроенных решения: argsort и argpartition ср. @Divakar ответ.

Однако, если производительность является фактором, то может (в зависимости от размеров данных и подмножества интереса) стоить противостоять "приманке одной строки", вложив еще одну строку и применив argsort к вывод 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 - это O (n log n), argpartition с аргументом диапазона - O (nk) (?), А argpartition + argsort - O (n + k log k)

Поэтому в интересном режиме n >> k >> 1 ожидается, что гибридный метод будет самым быстрым

ОБНОВЛЕНИЕ: 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)

Пример прогона:

top_k_sort 63.72379460372031
top_k_partition 99.30561298970133
top_k_hybrid 10.714635509066284