Keypoint Heatmap Argmax

Problem

Find the location of the maximum in a keypoint heatmap.

Return [row, col] of the largest value, keeping the first position in row-major order on ties.

case 1

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

expected: [0,1]

case 2

in: [[[7]]]

expected: [0,0]

case 3

in: [[[1,9,9],[3,9,2]]]

expected: [0,1]

Starter codePython

def heatmap_argmax(heatmap):
    # Your code here
    pass

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