Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

implementing argmax in Python

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)
like image 248
Neil G Avatar asked Sep 06 '25 12:09

Neil G


1 Answers

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)
like image 119
Neil G Avatar answered Sep 08 '25 13:09

Neil G