Argmax per Row

Problem

Return the index of the maximum entry in each row of A.

Break ties by returning the smallest index. Return a list of integers.

Hint

Scan left to right and only replace on a strict improvement to keep the first index.

case 1

in: [[[1,3,2],[5,0,4]]]

expected: [1,0]

case 2

in: [[[1,1],[2,2]]]

expected: [0,0]

case 3

in: [[[-5,-2,-9]]]

expected: [1]

case 4

in: [[[0]]]

expected: [0]

Starter codePython

def row_argmax(A):
    # Your code here
    pass

4 test cases · run your code against them in the editor