25

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)
Neil G
  • 32,138
  • 39
  • 156
  • 257
  • 2
    argmax on what? A function, a dictionary? – Swiss Feb 23 '11 at 23:42
  • on two iterables like in math: argmax_{keys} corresponding_values so that it returns the key corresponding to the greatest value. – Neil G Feb 23 '11 at 23:43
  • If they are associated key-value pairs, why not store them in an associative data structure such as a dictionary? – Swiss Feb 23 '11 at 23:47
  • 2
    This may be of interest to you as well: http://stackoverflow.com/questions/268272/getting-key-with-maximum-value-in-dictionary – Swiss Feb 23 '11 at 23:50
  • 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 Feb 23 '11 at 23:51
  • Cool, I didn't know that max could take a key. – Neil G Feb 23 '11 at 23:52
  • @Neil G: Please **update** the question so that it makes sense to someone who doesn't already know the answer. – S.Lott Feb 23 '11 at 23:54

5 Answers5

30

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)
Neil G
  • 32,138
  • 39
  • 156
  • 257
  • (I thought this was worth sharing since I couldn't find it on SO. If someone has a better solution, please share it!) – Neil G Feb 23 '11 at 23:24
  • I would recommend using regular zip instead of itertools.izip. Use of itertools is kind of frowned on unless you have a very good reason to use it. – Swiss Feb 24 '11 at 00:05
  • 1
    Although to be fair, izip returns an iterator while zip returns a list. So izip is faster if speed is important. – Swiss Feb 24 '11 at 00:11
  • 1
    You can also use lambda x: x[1] instead of operator.itemgetter(1). It looks a bit nicer, I think. – Swiss Feb 24 '11 at 00:15
  • @Swiss, as you point out izip is faster in some circumstances. Anyway, in python3, izip is renamed zip, right? – Neil G Feb 24 '11 at 00:15
  • 9
    @Swiss: I don't know where you got the idea that `itertools` was ever frowned upon. It was added to the standard library so it would be a readily available tool. – ncoghlan Feb 24 '11 at 02:52
  • @Neil: something that *is* frowned upon is using a `lambda` expression as the right hand side of an assignment statement. Use an ordinary named function instead. – ncoghlan Feb 24 '11 at 02:54
  • @ncoghlan, did not know that. I will change them now, but would you happen to have a reference by any chance? – Neil G Feb 24 '11 at 05:28
  • 2
    I don't know if it is written down anywhere - it's just in the assignment scenario you suffer all the drawbacks of using `lambda` without gaining any of the benefits, so a named function makes a lot more sense. – ncoghlan Feb 24 '11 at 05:48
  • @ncoghlan Usually there are more "pythonic" ways of doing the same things as what is provided by itertools. map and filter weren't even planned to be in Python 3 originally: http://www.artima.com/weblogs/viewpost.jsp?thread=98196 – Swiss Feb 25 '11 at 07:08
  • 2
    @Swiss you're confounding ``itertools`` with ``functools`` -- map and reduce are in the latter. – Elias Dorneles Nov 22 '14 at 18:50
21

Is the following code a fast and pythonic way?

idx_max = max(enumerate(x), key=lambda x:x[1])[0]
wall-e
  • 7,180
  • 2
  • 16
  • 9
  • 2
    While this is basically a stripped down and uncommented version of the accepted answer, I found it more useful, and a `lambda` is easier to read than `operator` stuff (even if you don't use it in a final solution) – loopbackbee Jul 23 '15 at 12:42
8
def argmax(lst):
     return lst.index(max(lst))

or analogously:

argmax = lambda lst: lst.index(max(lst)
Davoud Taghawi-Nejad
  • 16,142
  • 12
  • 62
  • 82
8

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.

Mikhail Korobov
  • 21,908
  • 8
  • 73
  • 65
6

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)
Andrew Clark
  • 202,379
  • 35
  • 273
  • 306
  • 1
    With list comprehension instead of imap: max([(func(*key), key) for key in keys]) – Swiss Feb 24 '11 at 00:23
  • 3
    Generator is pretty much identical if you want to match imap's speed. max((func(*key), key) for key in keys) – Swiss Feb 24 '11 at 00:33
  • @Swiss - list comprehension form is useful for those still on old versions of Python, pre-generator-expressions. But you are right, there is a performance penalty in creating the temporary list, just so you can get the max() of it. – PaulMcG Feb 24 '11 at 07:06