Início / coding / Perguntas / 78851114
Accepted
Mikael
Asked: 2024-08-09 11:49:45 +0800 CST

Indexação avançada Numpy em arrays multidimensionais; não apenas pegando linhas ou colunas simples

  • 772

Eu tenho uma matriz numpy tridimensional. Qual é a maneira mais rápida de obter um array 3D que tenha o maior item de cada eixo final do array sem escrever um loop. (Mais tarde usarei o CuPy com a mesma sintaxe, e os loops tirariam o paralelismo da GPU e a velocidade qual é o fator mais importante aqui.)

Obter os índices dos itens maiores é fácil:

>>> arr = np.array(
        [[[ 6, -2, -6, -5],
        [ 1, 12,  3,  9],
        [21,  7,  9,  8]],

       [[15, 12, 20, 12],
        [17, 15, 17, 23],
        [22, 18, 27, 32]]])
>>> indexes = arr.argmax(axis=2, keepdims=True)
>>> indexes
array([[[0],
        [1],
        [0]],

       [[2],
        [3],
        [3]]])

mas como usar esses índices para obter os valores escolhidos em arr? Todas as maneiras que tentei produziram erros (como arr[indexes]) ou resultados errados. O que eu gostaria de obter neste exemplo é

array([[[6],
        [12],
        [21]],

       [[20],
        [23],
        [32]]])
numpy

2 respostas

  1. Eu acho que você pode usar np.amax para isso

    import numpy as np

arr = np.array(
        [[[ 6, -2, -6, -5],
        [ 1, 12,  3,  9],
        [21,  7,  9,  8]],

       [[15, 12, 20, 12],
        [17, 15, 17, 23],
        [22, 18, 27, 32]]])

max_arr = np.amax(arr, axis=2, keepdims=True)
print(max_arr)

    Saída

    [[[ 6]
  [12]
  [21]]

 [[20]
  [23]
  [32]]]
    • 2
  2. Best Answer

    Conforme apontado por hpaulj, você pode usartake_along_axis

    >>>arr = np.array(
        [[[ 6, -2, -6, -5],
        [ 1, 12,  3,  9],
        [21,  7,  9,  8]],

       [[15, 12, 20, 12],
        [17, 15, 17, 23],
        [22, 18, 27, 32]]])
>>>indexes = arr.argmax(axis=2, keepdims=True)
>>>np.take_along_axis(arr , indexes , axis=2)

array([[[6],
        [12],
        [21]],

       [[20],
        [23],
        [32]]])
    • 2

relate perguntas