# Argmax in Python

Computing the argmax is a common programming problem: you have a set of candidates C, a score function f, and you want to compute an optimal candidate c in C according to f, i.e. one that maximizes f(c). If there is only one such candidate, it is called the "argmax" of f on C. One can find solutions for this on-line, e.g. in this post, but I realized recently that Python's max() function can actually handle the job. Here is how.

So, we have a set of candidates and we want to find the one that maximizes a score function, for instance:

In : candidates = [1, 5, 8, 11, 9]
In : f = lambda x: (x - 4) * (x - 10) * (x - 11)
In : [f(x) for x in candidates]
Out: [-270, 30, 24, 0, 10]


In this case the answer is 5, the candidate at which f is maximal with f(5)=30. A first solution could be to compute the values of f on candidates and use list.index() to return the index of a maximal element:

In : l = [f(x) for x in candidates]
In : fmax = max(l)
In : best_candidate = candidates[l.index(fmax)]
In : best_candidate
Out: 5


This approach relies on the fact that the indexes in l and candidates are the same, and it is a little bit verbose. There is a one-liner that circumvents this:

In : max(candidates, key=f)
Out: 5


In details: the key argument to max() is a one-argument function applied to the elements of the iterable before comparing them (its default value is of course the identity). For instance, max(i, j, key=f) is equivalent to i if f(i) >= f(j) else j.

In short:

argmax = lambda iterable, func: max(iterable, key=func)


## Discussion ¶

You can use Markdown with $\LaTeX$ formulas in your comment.