Como obtenho índices de N valores máximos em uma matriz NumPy?

484

NumPy propõe uma maneira de obter o índice do valor máximo de uma matriz via np.argmax.

Gostaria de algo semelhante, mas retornando os índices dos Nvalores máximos.

Por exemplo, se tiver uma matriz, [1, 3, 2, 4, 5], function(array, n=3)iria retornar os índices [4, 3, 1], que correspondem aos elementos [5, 4, 3].

Alexis Métaireau
fonte
4
Sua pergunta não está muito bem definida. Por exemplo, quais seriam os índices (você espera) array([5, 1, 5, 5, 2, 3, 2, 4, 1, 5]), com n= 3? Qual de todas as alternativas, como [0, 2, 3], [0, 2, 9], ...seria o correto? Por favor, elabore mais sobre seus requisitos específicos. Obrigado
comer
@eat, eu realmente não me importo com qual deles deve ser retornado neste caso específico. Mesmo que pareça lógico devolver o primeiro encontrado, isso não é um requisito para mim.
Alexis Métaireau
argsortpode ser uma alternativa viável se você não se importar com a ordem dos indeces retornados. Veja minha resposta abaixo.
blue

Respostas:

348

O mais simples que pude apresentar é:

In [1]: import numpy as np

In [2]: arr = np.array([1, 3, 2, 4, 5])

In [3]: arr.argsort()[-3:][::-1]
Out[3]: array([4, 3, 1])

Isso envolve uma classificação completa da matriz. Gostaria de saber se numpyfornece uma maneira interna de fazer uma classificação parcial; até agora não consegui encontrar um.

Se esta solução for muito lenta (especialmente para pequenas n ), pode valer a pena codificar algo no Cython .

NPE
fonte
1
A linha 3 poderia ser escrita equivalentemente como arr.argsort()[-1:-4:-1]? Eu tentei isso no intérprete e aparece com o mesmo resultado, mas estou me perguntando se não está quebrado por algum exemplo.
Abroekhof 20/09/12
44
@abroekhof Sim, isso deve ser equivalente a qualquer lista ou matriz. Como alternativa, isso poderia ser feito sem a reversão usando np.argsort(-arr)[:3], o que eu acho mais legível e direto ao ponto.
askewchan
6
o que significa [:: - 1]? @NPE
1a1a11a 17/10
@ 1a1a11a significa inverter uma matriz (literalmente, leva uma cópia de uma matriz a partir min sem restrições para max não constrangida na ordem inversa)
FizBack
15
arr.argsort()[::-1][:n]é melhor porque retorna vazio para em n=0vez de toda a matriz
abora
599

As versões mais recentes do NumPy (1.8 e superior) têm uma função chamada argpartitionpara isso. Para obter os índices dos quatro maiores elementos, faça

>>> a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])
>>> a[ind]
array([4, 9, 6, 9])

Diferentemente argsort, essa função é executada em tempo linear no pior dos casos, mas os índices retornados não são classificados, como pode ser visto no resultado da avaliação a[ind]. Se você precisar também, classifique-os depois:

>>> ind[np.argsort(a[ind])]
array([1, 8, 5, 0])

Obter os elementos top- k em ordem classificada dessa maneira leva tempo O ( n + k log k ).

Fred Foo
fonte
27
@varela argpartitioné executado em tempo linear, O (n), usando o algoritmo introselecionado . A classificação subsequente manipula apenas os elementos k, para que seja executado em O (k log k).
Fred Foo
2
Se alguém está se perguntando como exatamente np.argpartitione seu algoritmo irmã np.partitiontrabalho, há uma explicação mais detalhada na questão ligada: stackoverflow.com/questions/10337533/...
Ramon Martinez
7
@FredFoo: por que você usou -4? ?! você fez isso para começar para trás (desde k ser positivo ou negativo funciona da mesma para mim ele só imprime os menores números em primeiro lugar!
Rika
2
Uso @LKT a=np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])porque as listas normal do Python não suportam indexação por listas, ao contrárionp.array
Marawan Okasha
2
@Umangsinghal np.argpartitionusa um axisargumento opcional . Para encontrar os índices dos n principais valores para cada linha:np.argpartition(a, -n, axis=1)[-n:]
Ralph
48

Mais simples ainda:

idx = (-arr).argsort()[:n]

onde n é o número de valores máximos.

Ketan
fonte
7
Isso pode ser feito para uma matriz 2D? Se não, você sabe como?
Andrew Hundt
2
@AndrewHundt: basta usar (-arr) .argsort (axis = -1) [:,: n]
# 10
2
semelhante seria arr[arr.argsort()[-n:]]em vez de negar a matriz, basta levar uma fatia dos últimos n elementos
loganjones16
35

Usar:

>>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]

Para listas regulares de Python:

>>> a = [1, 3, 2, 4, 5]
>>> heapq.nlargest(3, range(len(a)), a.__getitem__)
[4, 3, 1]

Se você usa o Python 2, use em xrangevez de range.

Fonte: heapq - Algoritmo de fila da pilha

anishpatel
fonte
2
Não há necessidade de um loop em tudo aqui: heapq.nlargest(3, xrange(len(a)), a.take). Para listas Python, podemos usar em .__getitem__vez de .take.
Ashwini Chaudhary 28/10
Para matrizes n-dimensional Aem geral: heapq.nlargest(3, range(len(A.ravel())), A.ravel().take). (Espero que isso funcione apenas em ravel vs flatten
modos de exibição
31

Se estiver trabalhando com uma matriz multidimensional, será necessário achatar e desvendar os índices:

def largest_indices(ary, n):
    """Returns the n largest indices from a numpy array."""
    flat = ary.flatten()
    indices = np.argpartition(flat, -n)[-n:]
    indices = indices[np.argsort(-flat[indices])]
    return np.unravel_index(indices, ary.shape)

Por exemplo:

>>> xs = np.sin(np.arange(9)).reshape((3, 3))
>>> xs
array([[ 0.        ,  0.84147098,  0.90929743],
       [ 0.14112001, -0.7568025 , -0.95892427],
       [-0.2794155 ,  0.6569866 ,  0.98935825]])
>>> largest_indices(xs, 3)
(array([2, 0, 0]), array([2, 2, 1]))
>>> xs[largest_indices(xs, 3)]
array([ 0.98935825,  0.90929743,  0.84147098])
danvk
fonte
9

Se você não se importa com a ordem dos K-ésimas maiores elementos que pode usar argpartition, deve ter um desempenho melhor do que uma classificação completa argsort.

K = 4 # We want the indices of the four largest values
a = np.array([0, 8, 0, 4, 5, 8, 8, 0, 4, 2])
np.argpartition(a,-K)[-K:]
array([4, 1, 5, 6])

Os créditos vão para esta pergunta .

Fiz alguns testes e parece ter um argpartitiondesempenho superior à argsortmedida que o tamanho da matriz e o valor de K aumentam.

azul
fonte
7

Para matrizes multidimensionais, você pode usar a axispalavra - chave para aplicar o particionamento ao longo do eixo esperado.

# For a 2D array
indices = np.argpartition(arr, -N, axis=1)[:, -N:]

E para pegar os itens:

x = arr.shape[0]
arr[np.repeat(np.arange(x), N), indices.ravel()].reshape(x, N)

Mas observe que isso não retornará um resultado classificado. Nesse caso, você pode usar np.argsort()ao longo do eixo pretendido:

indices = np.argsort(arr, axis=1)[:, -N:]

# Result
x = arr.shape[0]
arr[np.repeat(np.arange(x), N), indices.ravel()].reshape(x, N)

Aqui está um exemplo:

In [42]: a = np.random.randint(0, 20, (10, 10))

In [44]: a
Out[44]:
array([[ 7, 11, 12,  0,  2,  3,  4, 10,  6, 10],
       [16, 16,  4,  3, 18,  5, 10,  4, 14,  9],
       [ 2,  9, 15, 12, 18,  3, 13, 11,  5, 10],
       [14,  0,  9, 11,  1,  4,  9, 19, 18, 12],
       [ 0, 10,  5, 15,  9, 18,  5,  2, 16, 19],
       [14, 19,  3, 11, 13, 11, 13, 11,  1, 14],
       [ 7, 15, 18,  6,  5, 13,  1,  7,  9, 19],
       [11, 17, 11, 16, 14,  3, 16,  1, 12, 19],
       [ 2,  4, 14,  8,  6,  9, 14,  9,  1,  5],
       [ 1, 10, 15,  0,  1,  9, 18,  2,  2, 12]])

In [45]: np.argpartition(a, np.argmin(a, axis=0))[:, 1:] # 1 is because the first item is the minimum one.
Out[45]:
array([[4, 5, 6, 8, 0, 7, 9, 1, 2],
       [2, 7, 5, 9, 6, 8, 1, 0, 4],
       [5, 8, 1, 9, 7, 3, 6, 2, 4],
       [4, 5, 2, 6, 3, 9, 0, 8, 7],
       [7, 2, 6, 4, 1, 3, 8, 5, 9],
       [2, 3, 5, 7, 6, 4, 0, 9, 1],
       [4, 3, 0, 7, 8, 5, 1, 2, 9],
       [5, 2, 0, 8, 4, 6, 3, 1, 9],
       [0, 1, 9, 4, 3, 7, 5, 2, 6],
       [0, 4, 7, 8, 5, 1, 9, 2, 6]])

In [46]: np.argpartition(a, np.argmin(a, axis=0))[:, -3:]
Out[46]:
array([[9, 1, 2],
       [1, 0, 4],
       [6, 2, 4],
       [0, 8, 7],
       [8, 5, 9],
       [0, 9, 1],
       [1, 2, 9],
       [3, 1, 9],
       [5, 2, 6],
       [9, 2, 6]])

In [89]: a[np.repeat(np.arange(x), 3), ind.ravel()].reshape(x, 3)
Out[89]:
array([[10, 11, 12],
       [16, 16, 18],
       [13, 15, 18],
       [14, 18, 19],
       [16, 18, 19],
       [14, 14, 19],
       [15, 18, 19],
       [16, 17, 19],
       [ 9, 14, 14],
       [12, 15, 18]])
Kasramvd
fonte
Eu acho que você pode simplificar a indexação aqui usando np.take_along_axis(o que provavelmente não existia quando você respondeu a esta pergunta)
Eric
4

Isso será mais rápido que uma classificação completa, dependendo do tamanho da sua matriz original e do tamanho da sua seleção:

>>> A = np.random.randint(0,10,10)
>>> A
array([5, 1, 5, 5, 2, 3, 2, 4, 1, 0])
>>> B = np.zeros(3, int)
>>> for i in xrange(3):
...     idx = np.argmax(A)
...     B[i]=idx; A[idx]=0 #something smaller than A.min()
...     
>>> B
array([0, 2, 3])

Evidentemente, isso envolve adulterar sua matriz original. Que você pode corrigir (se necessário) fazendo uma cópia ou substituindo novamente os valores originais. ... o que for mais barato para o seu caso de uso.

Paulo
fonte
FWIW, sua solução não fornecerá uma solução inequívoca em todas as situações. O OP deve descrever como lidar com esses casos inequívocos. Obrigado
comer
@eat A pergunta do OP é um pouco ambígua. Uma implementação, no entanto, não é realmente aberta à interpretação. :) O OP deve se referir simplesmente à definição de np.argmax docs.scipy.org/doc/numpy/reference/generated/numpy.argmax.html para garantir que essa solução específica atenda aos requisitos. É possível que qualquer solução reunião reqirement declarado do OP é aceitável ..
Paul
Bem, pode-se considerar a implementação de argmax(.)também não ambígua. (IMHO tenta seguir algum tipo de lógica de curto-circuito, mas infelizmente falha em fornecer um comportamento universalmente aceitável). Obrigado
comer
3

O método np.argpartitionretorna apenas os k maiores índices, executa uma classificação local e é mais rápido que np.argsort(executando uma classificação completa) quando a matriz é bastante grande. Mas os índices retornados NÃO estão em ordem crescente / decrescente . Digamos com um exemplo:

Digite a descrição da imagem aqui

Podemos ver que, se você deseja uma ordem ascendente estrita dos principais k índices, np.argpartition não retornará o que deseja.

Além de fazer uma classificação manualmente após o np.argpartition, minha solução é usar o PyTorch, torch.topk uma ferramenta para construção de redes neurais, fornecendo APIs do tipo NumPy com suporte a CPU e GPU. É tão rápido quanto o NumPy com MKL e oferece um aumento de GPU se você precisar de grandes cálculos de matriz / vetor.

O código estrito de subida / descida dos índices principais k será:

Digite a descrição da imagem aqui

Observe que torch.topkaceita um tensor da tocha e retorna os valores de k top e os índices de k top no tipo torch.Tensor. Semelhante ao np, o torch.topk também aceita um argumento de eixo para que você possa manipular matrizes / tensores multidimensionais.

futuro
fonte
2

Usar:

from operator import itemgetter
from heapq import nlargest
result = nlargest(N, enumerate(your_list), itemgetter(1))

Agora a resultlista conterá N tuplas ( index, value) onde valueé maximizada.

off99555
fonte
2

Usar:

def max_indices(arr, k):
    '''
    Returns the indices of the k first largest elements of arr
    (in descending order in values)
    '''
    assert k <= arr.size, 'k should be smaller or equal to the array size'
    arr_ = arr.astype(float)  # make a copy of arr
    max_idxs = []
    for _ in range(k):
        max_element = np.max(arr_)
        if np.isinf(max_element):
            break
        else:
            idx = np.where(arr_ == max_element)
        max_idxs.append(idx)
        arr_[idx] = -np.inf
    return max_idxs

Também funciona com matrizes 2D. Por exemplo,

In [0]: A = np.array([[ 0.51845014,  0.72528114],
                     [ 0.88421561,  0.18798661],
                     [ 0.89832036,  0.19448609],
                     [ 0.89832036,  0.19448609]])
In [1]: max_indices(A, 8)
Out[1]:
    [(array([2, 3], dtype=int64), array([0, 0], dtype=int64)),
     (array([1], dtype=int64), array([0], dtype=int64)),
     (array([0], dtype=int64), array([1], dtype=int64)),
     (array([0], dtype=int64), array([0], dtype=int64)),
     (array([2, 3], dtype=int64), array([1, 1], dtype=int64)),
     (array([1], dtype=int64), array([1], dtype=int64))]

In [2]: A[max_indices(A, 8)[0]][0]
Out[2]: array([ 0.89832036])
X Æ A-12
fonte
Funciona bem, mas fornece mais resultados se você tiver valores duplicados (máximos) em sua matriz A. Eu esperaria exatamente k resultados, mas no caso de valores duplicados, você obtém mais de k resultados.
Guido
Eu modifiquei levemente o código. A lista de índices retornados tem comprimento igual exatamente a k. Se você tiver duplicatas, elas serão agrupadas em uma única tupla.
X Æ A-12
1

bottleneck possui uma função de classificação parcial, se a despesa de classificar toda a matriz apenas para obter os N maiores valores for muito grande.

Não sei nada sobre este módulo; Eu apenas pesquisei no Google numpy partial sort.

Katriel
fonte
Eu não encontro nenhuma função de classificação parcial no gargalo, há uma função de partição, mas isso não faz tipo
nbecker
1

A seguir, é uma maneira muito fácil de ver o máximo de elementos e suas posições. Aqui axisestá o domínio; axis= 0 significa o número máximo em colunas e axis= 1 significa o número máximo em linhas para o caso 2D. E para dimensões mais altas depende de você.

M = np.random.random((3, 4))
print(M)
print(M.max(axis=1), M.argmax(axis=1))
liberal
fonte
0

Eu achei mais intuitivo de usar np.unique.

A idéia é que o método exclusivo retorne os índices dos valores de entrada. Então, a partir do valor único máximo e das indicações, a posição dos valores originais pode ser recriada.

multi_max = [1,1,2,2,4,0,0,4]
uniques, idx = np.unique(multi_max, return_inverse=True)
print np.squeeze(np.argwhere(idx == np.argmax(uniques)))
>> [4 7]
phi
fonte
0

Eu acho que a maneira mais eficiente em termos de tempo é iterar manualmente através da matriz e manter uma pilha mínima de tamanho k, como outras pessoas mencionaram.

E também proponho uma abordagem de força bruta:

top_k_index_list = [ ]
for i in range(k):
    top_k_index_list.append(np.argmax(my_array))
    my_array[top_k_index_list[-1]] = -float('inf')

Defina o maior elemento para um grande valor negativo depois de usar argmax para obter seu índice. E então a próxima chamada de argmax retornará o segundo maior elemento. E você pode registrar o valor original desses elementos e recuperá-los, se desejar.

Zhenghao Zhao
fonte
0

Este código funciona para uma matriz matricial numpy:

mat = np.array([[1, 3], [2, 5]]) # numpy matrix

n = 2  # n
n_largest_mat = np.sort(mat, axis=None)[-n:] # n_largest 
tf_n_largest = np.zeros((2,2), dtype=bool) # all false matrix
for x in n_largest_mat: 
  tf_n_largest = (tf_n_largest) | (mat == x) # true-false  

n_largest_elems = mat[tf_n_largest] # true-false indexing 

Isso produz uma indexação de matriz n_largest true-false que também funciona para extrair n_largest elementos de uma matriz de matriz

Yi Xiang Chong
fonte