How can I take a 2d array and then take every row and collapse it into the index of the maximum numb-CodePudding

I have a 2d array and I am trying to get the index and value of the largest number on the second axis. For example:

[
[1, 5, 3],
[6, 2, 4],
[4, 3, 5]
]

would return [indexofmax, max] for every line.

So:

[
[1, 5],
[0, 6],
[2, 5]
]

CodePudding user response：

This line comprehension should do it:

def get_max_index(arr):
    return [[subarr.index(max(subarr)), max(subarr)] for subarr in arr]

Note - This will only return the first index of the largest number.

CodePudding user response：

Since you tagged numpy, try with argmax and take_along axis:

# convert to numppy array if not already is
arr = np.array(arr)

idx = np.argmax(arr, axis=1)[..., None]     # the index of row maxmimum
np.hstack([idx, np.take_along_axis(arr, idx, axis=1)])