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]

Published by

Daniel Lemire

A computer science professor at the Université du Québec (TELUQ).

10 thoughts on “Computing argmax fast in Python”

  1. 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]

  2. You should also try a functional idiom – typically better and faster:

    array = [1.0, 2.0, 3.0, 1.0]

    m = reduce(max, array)

  3. 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!

  4. 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.

  5. Pingback: argmax in Python | Tastalian's Blog

Leave a Reply

Your email address will not be published. Required fields are marked *

To create code blocks or other preformatted text, indent by four spaces:

    This will be displayed in a monospaced font. The first four 
    spaces will be stripped off, but all other whitespace
    will be preserved.
    
    Markdown is turned off in code blocks:
     [This is not a link](http://example.com)

To create not a block, but an inline code span, use backticks:

Here is some inline `code`.

For more help see http://daringfireball.net/projects/markdown/syntax