Computing argmax fast in Python
Update: see Fast argmax in Python for the final word.
Python doesn’t come with an argmax function built-in. That is, a function that tells you where a maximum is in an array. And I keep needing to write my own argmax function so I’m getting better and better at it. Here’s the best I could come up with so far:
from itertools import izip
argmax = lambda array: max(izip(array, xrange(len(array))))[1]
The really nice thing about izip and xrange is that they don’t actually output arrays, but only lazy iterators. You also have plenty of similar functions such as imap or ifilter. Very neat.
Here’s a challenge to you: can you do better? That is, can you write a faster, less memory hungry function? If not, did I find the optimal solution? If so, do I get some kind of prize?
Next, suppose you want to find argmax, but excluding some set of bad indexes, you can do it this way…
from itertools import izip
argmaxbad = lambda array,badindexes: max(ifilter(lambda t: t[1] not in badindexes,izip(array, xrange(len(array)))))[1]
Python is amazingly easy.
As a side-note, you can also do intersections and union of sets in Python in a similar (functional) spirit:
def intersection(set1,set2): return filter(lambda s:s in set2,set1)
def union(set1, set2): return set1 + filter(lambda s:s not in set1, set2)
Update: you can do the same things with hash tables:
max(izip(hashtable[max].itervalues(), hashtable[max].iterkeys()))[1]
Montreal, Canada 
Follow on
You can try this:
def positionMax(sequence, excludedIndexesSet=None):
if not badindexes:
return max( izip(sequence, xrange(len(sequence)) ) )[1]
else:
badindexes = set(badindexes)
return max( (e,i) for i,e in enumerate(sequence) if i not in badindexes )[1]
Comment by Anonymous — 13/2/2005 @ 9:16
You should also try a functional idiom – typically better and faster:
array = [1.0, 2.0, 3.0, 1.0]
m = reduce(max, array)
Comment by SwiftCoder — 29/1/2008 @ 9:20
I’m sorry SwiftCoder, but I believe that is a max. It is equivalent to:
m = max(array)
My (arg)max:
lambda array: max([array[i],i] for i in range(len(array)))[1]
Cheers!
Comment by Protector one — 7/5/2008 @ 15:31
array.index(max(array))
is it too simple?
Comment by kralin — 17/12/2008 @ 6:27
[...] my post Computing argmax fast in Python, I reported that Python has no builtin function to compute argmax, the position of a maximal value. [...]
Pingback by Fast argmax in Python — 17/12/2008 @ 8:48
import numpy
array = numpy.array ( [1,2,4,3,5,1 ] )
numpy.argmax( array )
# should return 4
Comment by Sam — 18/12/2008 @ 8:48
In case of duplicate values you solution returns the last position, while the array.index(max(array)) solution returns the first which, beside the simplicity of the solution, seems to be the more natural answer.
Comment by Mapio — 29/1/2011 @ 11:41
How about
max(xrange(len(array)), key=array.__getitem__)
Comment by Anonymous — 27/4/2011 @ 10:59
[...] called the “argmax” of (f) on (C). One can find solutions for this on-line, e.g. in this blog post, but I realized recently that Python’s max() function can actually handle the job. Here is [...]
Pingback by argmax in Python | Tastalian's Blog — 24/4/2012 @ 17:55