28

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)
Improve this question
7
  • 2
    argmax on what? A function, a dictionary? Commented Feb 23, 2011 at 23:42
  • on two iterables like in math: argmax_{keys} corresponding_values so that it returns the key corresponding to the greatest value. Commented Feb 23, 2011 at 23:43
  • If they are associated key-value pairs, why not store them in an associative data structure such as a dictionary? Commented Feb 23, 2011 at 23:47
  • 2
    This may be of interest to you as well: stackoverflow.com/questions/268272/… Commented Feb 23, 2011 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). Commented Feb 23, 2011 at 23:51

5 Answers 5

Reset to default
34

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)
Improve this answer
Sign up to request clarification or add additional context in comments.

11 Comments

Although to be fair, izip returns an iterator while zip returns a list. So izip is faster if speed is important.
You can also use lambda x: x[1] instead of operator.itemgetter(1). It looks a bit nicer, I think.
@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.
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.
@Swiss you're confounding itertools with functools -- map and reduce are in the latter.
|
22

Is the following code a fast and pythonic way?

idx_max = max(enumerate(x), key=lambda x:x[1])[0]
Improve this answer

1 Comment

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)
10 
def argmax(lst):
     return lst.index(max(lst))

or analogously:

argmax = lambda lst: lst.index(max(lst)
Improve this answer

2 Comments

This approach is theoretically suboptimal because it iterates over the list twice. However, it is elegant.
In fact, this one is faster then any other solutions posted here.
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.

Improve this answer

Comments

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)
Improve this answer

3 Comments

With list comprehension instead of imap: max([(func(*key), key) for key in keys])
Generator is pretty much identical if you want to match imap's speed. max((func(*key), key) for key in keys)
@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.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.