Repository: more-itertools/more-itertools

Issue:
argmin() and argmax()
These two functions are common in machine learning.  

[Numpy implements them](https://numpy.org/doc/2.2/reference/generated/numpy.argmin.html) as returning the index of the first occurrence of a minimum or maximum value in an array. For more_itertools, the functions would accept any iterable or iterator and make only a single pass over the data one value at a time.


```
from operator import itemgetter

def argmin(iterable, *, key=None):
    "Index of the first occurence of a minimum value in an iterable."
    if key is not None:
        iterable = map(key, iterable)
    return min(enumerate(iterable), key=itemgetter(1))[0] 

def argmax(iterable, *, key=None):
    "Index of the first occurence of a maximum value in an iterable."
    if key is not None:
        iterable = map(key, iterable)
    return max(enumerate(iterable), key=itemgetter(1))[0] 
```

Some tests:

```
data = [10, 2, 20, 5, 17, 4]
assert argmin(data) == 1
assert argmax(data) == 2

data = [10, -2, -20, 5, 17, 4]
assert argmin(data, key=abs) == 1
assert argmax(data, key=abs) == 2

data = [[0]*10, [0]*5, [0]*3, [0]*12, [0]*2, [0]*3]
assert argmin(data, key=len) == 4
assert argmax(data, key=len) == 3

assert argmin([10, 10, 20, 10]) == 0
assert argmax([30, 30, 20, 30]) == 0
```

Discussion:
@bbayles: Nice and compact implementations. I fee like as soon as we have this, the question will be "what about a function that returns both the index and the actual value?" So should the `[0]` be left out in the `return` statements?@rhettinger: I have mixed feelings about this. 

People likely would ask for this. Since the input can be an iterator, there might not be an easy way to go from the index back to the returned value.

OTOH the standard math function returns only the argument to the function (which is why it is called ArgMin and ArgMax). Also, functions that return two values are awkward to use. They don't nest well inside larger expressions, and it is no fun to write `i, _ = argmin(data, key=cost)` or `i = argmin(data, key=cost)[0]`.@rhettinger: I suggest sticking with the standard math function as a good minimum viable product. 

If a real use arises for having both the value and index (which i doubt), it wouldn't difficult to later add a second function or to add a switch to the existing function: `argmin(iterable, *, key=None, index_value=False)`

The reason for my doubt is that possible return values do not naturally co-occur in problems. We can be interested in the value from the iterable, the result of applying a key function to that value, or to its position.  And the only reason to get the position is to index into some other feature array.  So ISTM that the position and the value are usually orthogonal needs because they answer completely different questions.

In the following example, which minimizes a cost function (the time for a wound to head), the only reason to get the position is to locate a corresponding label.

```python

# Time for a wound to heal
cost = lambda x: x**2 - 20*x + 150

indices = [   0,       1,       2,      3,      4    ]
labels =  ['homer', 'marge', 'bart', 'lisa', 'maggie']
ages =    [  35,      30,      10,      9,      1    ]

# Which participant age heals the fastest?
>>> min(ages, key=cost)
10

# Least number of days until a scar heals?
>>> cost(_)
50

# Name the fastest healing family member.
>>> labels[argmin(ages, key=cost)]
'bart'
```

I like that these are one-liners that answer the question directly without having to add a `[0]` or `[1]` in the expression.@bbayles: I think I'd have to relent and add secondary functions that return both, because I think complicating the return type based on an input keyword argument would be too weird.

I'll add these two as-is for now.

Solution patch:
diff --git a/docs/api.rst b/docs/api.rst
index 5b560c2..6753c28 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -153,6 +153,8 @@ These tools return summarized or aggregated data from an iterable.
 .. autofunction:: exactly_n(iterable, n, predicate=bool)
 .. autofunction:: is_sorted
 .. autofunction:: all_unique
+.. autofunction:: argmin
+.. autofunction:: argmax
 .. function:: minmax(iterable, *[, key, default])
 .. autofunction:: minmax(arg1, arg2, *args[, key])
   :noindex:
diff --git a/more_itertools/more.py b/more_itertools/more.py
index 1bf48ba..ae3ab92 100755
--- a/more_itertools/more.py
+++ b/more_itertools/more.py
@@ -53,6 +53,8 @@ __all__ = [
     'all_unique',
     'always_iterable',
     'always_reversible',
+    'argmax',
+    'argmin',
     'bucket',
     'callback_iter',
     'chunked',
@@ -5067,3 +5069,31 @@ def nth_prime(n):
         raise ValueError
     limit = math.ceil(_nth_prime_ub(n + 1))
     return nth(sieve(limit), n)
+
+
+def argmin(iterable, *, key=None):
+    """
+    Index of the first occurence of a minimum value in an iterable.
+
+    >>> argmin('efghabcdijkl')
+    4
+    >>> argmin([3, 2, 1, 0, 4, 2, 1, 0])
+    3
+    """
+    if key is not None:
+        iterable = map(key, iterable)
+    return min(enumerate(iterable), key=itemgetter(1))[0]
+
+
+def argmax(iterable, *, key=None):
+    """
+    Index of the first occurence of a maximum value in an iterable.
+
+    >>> argmax('abcdefghabcd')
+    7
+    >>> argmax([0, 1, 2, 3, 3, 2, 1, 0])
+    3
+    """
+    if key is not None:
+        iterable = map(key, iterable)
+    return max(enumerate(iterable), key=itemgetter(1))[0]
diff --git a/more_itertools/more.pyi b/more_itertools/more.pyi
index f06b35d..bce9bd8 100644
--- a/more_itertools/more.pyi
+++ b/more_itertools/more.pyi
@@ -34,6 +34,8 @@ __all__ = [
     'all_unique',
     'always_iterable',
     'always_reversible',
+    'argmax',
+    'argmin',
     'bucket',
     'callback_iter',
     'chunked',
@@ -933,3 +935,9 @@ def dft(xarr: Sequence[complex]) -> Iterator[complex]: ...
 def idft(Xarr: Sequence[complex]) -> Iterator[complex]: ...
 def _nth_prime_ub(n: int) -> float: ...
 def nth_prime(n: int) -> int: ...
+def argmin(
+    iterable: Iterable[_T], *, key: Callable[[_T], _U] | None = ...
+) -> int: ...
+def argmax(
+    iterable: Iterable[_T], *, key: Callable[[_T], _U] | None = ...
+) -> int: ...


Test patch:
diff --git a/tests/test_more.py b/tests/test_more.py
index 5e45f7e..d3717b5 100644
--- a/tests/test_more.py
+++ b/tests/test_more.py
@@ -6051,3 +6051,31 @@ class DoubleStarMapTests(TestCase):
         actual = list(mi.doublestarmap(lambda x: x, []))
         expected = []
         self.assertEqual(actual, expected)
+
+
+class ArgMinArgMaxTests(TestCase):
+    def test_basic(self):
+        for i, iterable, expected_min, expected_max in (
+            (1, [10, 2, 20, 5, 17, 4], 1, 2),
+            (2, [10, -2, -20, 5, 17, 4], 2, 4),
+            (3, [10, 10, 20, 10], 0, 2),
+            (4, [30, 30, 20, 30], 2, 0),
+        ):
+            with self.subTest(i=i):
+                self.assertEqual(mi.argmin(iterable), expected_min)
+                self.assertEqual(mi.argmax(iterable), expected_max)
+
+    def test_key(self):
+        for i, iterable, key, expected_min, expected_max in (
+            (1, [10, -2, -20, 5, 17, 4], abs, 1, 2),
+            (
+                2,
+                [[0] * 10, [0] * 5, [0] * 3, [0] * 12, [0] * 2, [0] * 3],
+                len,
+                4,
+                3,
+            ),
+        ):
+            with self.subTest(i=i):
+                self.assertEqual(mi.argmin(iterable, key=key), expected_min)
+                self.assertEqual(mi.argmax(iterable, key=key), expected_max)