如何返回二維 NumPy 數組中最高和元素的索引？

Question

假設我有以下數組：

array = numpy.array([[[7.1, 4.5, 2.1], [0.9, 0.6, 10.1]],
                    [[11.0, 5.4, 4.3], [6.7, 0.3, 8.2]]])

我正在嘗試創建一個函數，該函數允許我找到總和值最高的元素並以形式（行、列）返回其位置，可能作為元組。 在這種情況下，所需的輸出將是(1,0)作為第二行中的元素，第一列的總和值最高。

這是我的方法，但我似乎無法弄清楚如何獲得我想要的輸出。 我已經廣泛瀏覽了 NumPy 文檔，並嘗試了 .max 和 .where 命令的許多變體。

def function(array):
for row in range(len(array)):
    for column in range(len(array[row])):
        total = (array[row, column, 0] + array[row, column, 1] + array[row, column, 2]) 
        array[row, column, 0] = total
        array[row, column, 1] = total
        array[row, column, 2] = total
    return ???

我的想法是將總值重新分配給數組中每個列表中的所有元素，然后使用 numpy.max 之類的東西給我最大值的位置，但這似乎產生了超出范圍的索引值期待。

任何建議表示贊賞！

Answer 1

使用np.argmax查找最大索引並使用np.unravel_index將其轉換為預期格式：

import numpy as np

array = np.array([[[7.1, 4.5, 2.1], [0.9, 0.6, 10.1]],
                  [[11.0, 5.4, 4.3], [6.7, 0.3, 8.2]]])

# sum across last index to find values
ssum = array.sum(-1)

# find the index using argmax, un-ravel it
res = np.unravel_index(ssum.argmax(), ssum.shape)
print(res)

輸出

(1, 0)

Answer 2

您需要先計算總和，然后找到最大值並找到索引，如下所示：

>>> import numpy as np
>>> arr = np.array([[[7.1, 4.5, 2.1], [0.9, 0.6, 10.1]],
...                    [[11.0, 5.4, 4.3], [6.7, 0.3, 8.2]]])

>>> sum_3d = np.sum(arr,axis = 2)
>>> sum_3d
array([[13.7, 11.6],
       [20.7, 15.2]])

>>> max_sum_3d = np.amax(sum_3d)
>>> max_sum_3d
20.7

>>> np.argwhere( sum_3d == max_sum_3d )
array([[1, 0]])

# If you want output as tuple you can try this
>>> idx = np.argwhere( sum_3d == max_sum_3d )
>>> tuple(idx[0])
(1,0)