How to return index of highest sum element in 2D NumPy array?-CodePudding

Suppose I have the following array:

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

I am trying to create a function that allows me to find the element with the highest sum value and return its position in the form (row, column), probably as a tuple. In this case, the desired output would be (1,0) as the element in the second row, first column has the highest sum value.

This has been my approach, but I can't seem to figure out how to get the output I desire. I've looked through the NumPy documentation extensively and tried many variations of .max and .where commands.

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 ???

My thought was to reassign the total value to all the elements within each list in the array then use something like numpy.max to give me the position of the maximum value, but that appears to produce an index value outside of the range I'm expecting.

Any advice is appreciated!

CodePudding user response：

Use np.argmax to find the maximum index and np.unravel_index to convert it to the expected format:

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)

Output

(1, 0)

CodePudding user response：

You need first compute sum then find max and find index like below:

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