Subgradient of Max

Problem

The subgradient of max_i x_i is a one-hot vector selecting an index that attains the maximum.

Given a list of values, return the one-hot vector choosing the first index with the largest value, which is a valid subgradient on ties.

Hint

At a tie any convex combination of the maximizers is a valid subgradient.

case 1

in: [[1,3,2]]

expected: [0,1,0]

case 2

in: [[5,5,1]]

expected: [1,0,0]

case 3

in: [[-1,-2]]

expected: [1,0]

case 4

in: [[0]]

expected: [1]

Starter codePython

def subgradient_max(values):
    # Your code here
    pass

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