Python + numpy: эффективный способ получения значений min/max n и индексов из матрицы

Какой эффективный способ, учитывая матрицу numpy (массив 2-d), вернуть значения min/max n (вместе с их индексами) в массиве? В настоящее время у меня есть:

def n_max(arr, n):
    res = [(0,(0,0))]*n
    for y in xrange(len(arr)):
        for x in xrange(len(arr[y])):
            val = float(arr[y,x])
            el = (val,(y,x))
            i = bisect.bisect(res, el)
            if i > 0:
                res.insert(i, el)
                del res[0]
    return res

Это занимает 3 раза дольше, чем алгоритм сопоставления шаблонов изображений, который pyopencv делает, чтобы сгенерировать массив, в котором я хочу запустить это, и я считаю, что это глупо.

Ответ 1

Со времени другого ответа NumPy добавил numpy.partition и numpy.argpartition для частичной сортировки, что позволяет сделать это в O(arr.size) time или O(arr.size+n*log(n)), если вам нужны элементы в отсортированном порядке.

numpy.partition(arr, n) возвращает массив размером arr, где элемент n th - это то, что было бы, если массив был отсортирован. Все меньшие элементы появляются перед этим элементом, и все более крупные элементы появляются позже.

numpy.argpartition соответствует numpy.partition, поскольку numpy.argsort соответствует numpy.sort.

Здесь вы можете использовать эти функции для поиска индексов минимальных элементов n arr:

flat_indices = numpy.argpartition(arr.ravel(), n-1)[:n]
row_indices, col_indices = numpy.unravel_index(flat_indices, arr.shape)

И если вам нужны индексы по порядку, то row_indices[0] - это строка минимального элемента, а не только один из минимальных элементов n:

min_elements = arr[row_indices, col_indices]
min_elements_order = numpy.argsort(min_elements)
row_indices, col_indices = row_indices[min_elements_order], col_indices[min_elements_order]

Ответ 2

Так как в NumPy нет реализации кучи, возможно, лучше всего вы хотите отсортировать весь массив и взять последние n элементы:

def n_max(arr, n):
    indices = arr.ravel().argsort()[-n:]
    indices = (numpy.unravel_index(i, arr.shape) for i in indices)
    return [(arr[i], i) for i in indices]

(Это, вероятно, вернет список в обратном порядке по сравнению с вашей реализацией - не проверял.)

Изменить. Более эффективное решение, которое работает с более новыми версиями Numpy, приведено в этом ответе