Skip to content

Commit

Permalink
PERF: RangeIndex.argmin/argmax (pandas-dev#57823)
Browse files Browse the repository at this point in the history
  • Loading branch information
mroeschke authored Mar 14, 2024
1 parent 97c31a6 commit d79910c
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 1 deletion.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v3.0.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,7 @@ Performance improvements
- Performance improvement in :meth:`MultiIndex.equals` for equal length indexes (:issue:`56990`)
- Performance improvement in :meth:`RangeIndex.__getitem__` with a boolean mask or integers returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57588`)
- Performance improvement in :meth:`RangeIndex.append` when appending the same index (:issue:`57252`)
- Performance improvement in :meth:`RangeIndex.argmin` and :meth:`RangeIndex.argmax` (:issue:`57823`)
- Performance improvement in :meth:`RangeIndex.round` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57824`)
- Performance improvement in :meth:`RangeIndex.join` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57651`, :issue:`57752`)
- Performance improvement in :meth:`RangeIndex.reindex` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57647`, :issue:`57752`)
Expand Down
35 changes: 34 additions & 1 deletion pandas/core/indexes/range.py
Original file line number Diff line number Diff line change
Expand Up @@ -515,7 +515,7 @@ def copy(self, name: Hashable | None = None, deep: bool = False) -> Self:
new_index = self._rename(name=name)
return new_index

def _minmax(self, meth: str) -> int | float:
def _minmax(self, meth: Literal["min", "max"]) -> int | float:
no_steps = len(self) - 1
if no_steps == -1:
return np.nan
Expand All @@ -536,6 +536,39 @@ def max(self, axis=None, skipna: bool = True, *args, **kwargs) -> int | float:
nv.validate_max(args, kwargs)
return self._minmax("max")

def _argminmax(
self,
meth: Literal["min", "max"],
axis=None,
skipna: bool = True,
) -> int:
nv.validate_minmax_axis(axis)
if len(self) == 0:
return getattr(super(), f"arg{meth}")(
axis=axis,
skipna=skipna,
)
elif meth == "min":
if self.step > 0:
return 0
else:
return len(self) - 1
elif meth == "max":
if self.step > 0:
return len(self) - 1
else:
return 0
else:
raise ValueError(f"{meth=} must be max or min")

def argmin(self, axis=None, skipna: bool = True, *args, **kwargs) -> int:
nv.validate_argmin(args, kwargs)
return self._argminmax("min", axis=axis, skipna=skipna)

def argmax(self, axis=None, skipna: bool = True, *args, **kwargs) -> int:
nv.validate_argmax(args, kwargs)
return self._argminmax("max", axis=axis, skipna=skipna)

def argsort(self, *args, **kwargs) -> npt.NDArray[np.intp]:
"""
Returns the indices that would sort the index and its
Expand Down
24 changes: 24 additions & 0 deletions pandas/tests/indexes/ranges/test_range.py
Original file line number Diff line number Diff line change
Expand Up @@ -763,6 +763,30 @@ def test_getitem_boolmask_wrong_length():
ri[[True]]


@pytest.mark.parametrize(
"rng",
[
range(0, 5, 1),
range(0, 5, 2),
range(10, 15, 1),
range(10, 5, -1),
range(10, 5, -2),
range(5, 0, -1),
],
)
@pytest.mark.parametrize("meth", ["argmax", "argmin"])
def test_arg_min_max(rng, meth):
ri = RangeIndex(rng)
idx = Index(list(rng))
assert getattr(ri, meth)() == getattr(idx, meth)()


@pytest.mark.parametrize("meth", ["argmin", "argmax"])
def test_empty_argmin_argmax_raises(meth):
with pytest.raises(ValueError, match=f"attempt to get {meth} of an empty sequence"):
getattr(RangeIndex(0), meth)()


def test_getitem_integers_return_rangeindex():
result = RangeIndex(0, 10, 2, name="foo")[[0, -1]]
expected = RangeIndex(start=0, stop=16, step=8, name="foo")
Expand Down

0 comments on commit d79910c

Please sign in to comment.