itertools - implementing argmax in Python

25
votes

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)
5
pythonitertools
argmax on what? A function, a dictionary? - Swiss
on two iterables like in math: argmax_{keys} corresponding_values so that it returns the key corresponding to the greatest value. - Neil G
If they are associated key-value pairs, why not store them in an associative data structure such as a dictionary? - Swiss
This may be of interest to you as well: stackoverflow.com/questions/268272/… - Swiss
I guess you could do that, but my solution to argmax would have to unzip dict.items() and them re-zip it as (values, keys). - Neil G

5 Answers

30
votes

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)
21
votes

Is the following code a fast and pythonic way?

idx_max = max(enumerate(x), key=lambda x:x[1])[0]
8
votes

I found this way easier to think about argmax: say we want to calculate argmax(f(y)) where y is an item from Y. So for each y we want to calculate f(y) and get y with maximum f(y).

This definition of argmax is general, unlike "given an iterable of values return the index of the greatest value" (and it is also quite natural IMHO).

And ..drumroll.. Python allows to do exactly this using a built-in max:

best_y = max(Y, key=f)

So argmax_f (from the accepted answer) is unnecesary complicated and inefficient IMHO - it is a complicated version of built-in max. All other argmax-like tasks should become clear at this point: just define a proper function f.

7
votes
def argmax(lst):
     return lst.index(max(lst))

or analogously:

argmax = lambda lst: lst.index(max(lst)
6
votes

Based on Neil's answer, but specialized for functions that take multiple arguments.

argmax = lambda keys, func: max(imap(lambda key: (func(*key), key), keys))[1]

For example:

argmax([(5, 2), (3, 3), (2, 5)], pow)
# (2, 5)