implementing argmax in Python

Neil G picture Neil G · Feb 24, 2011 · Viewed 15.8k times · Source

How should argmax be implemented in Python? It should be as efficient as possible, so it should work with iterables.

Three ways it could be implemented:

  • given an iterable of pairs return the key corresponding to the greatest value
  • given an iterable of values return the index of the greatest value
  • given an iterable of keys and a function f, return the key with largest f(key)

Answer

Neil G picture Neil G · Feb 24, 2011

I modified the best solution I found:

# given an iterable of pairs return the key corresponding to the greatest value
def argmax(pairs):
    return max(pairs, key=lambda x: x[1])[0]

# given an iterable of values return the index of the greatest value
def argmax_index(values):
    return argmax(enumerate(values))

# given an iterable of keys and a function f, return the key with largest f(key)
def argmax_f(keys, f):
    return max(keys, key=f)