|
64 | 64 | _dtype_int64 = np.dtype(np.int64)
|
65 | 65 |
|
66 | 66 |
|
| 67 | +def min_fitting_element(start: int, step: int, lower_limit: int) -> int: |
| 68 | + """Returns the smallest element greater than or equal to the limit""" |
| 69 | + no_steps = -(-(lower_limit - start) // abs(step)) |
| 70 | + return start + abs(step) * no_steps |
| 71 | + |
| 72 | + |
67 | 73 | class RangeIndex(Index):
|
68 | 74 | """
|
69 | 75 | Immutable Index implementing a monotonic integer range.
|
@@ -570,25 +576,30 @@ def argsort(self, *args, **kwargs) -> npt.NDArray[np.intp]:
|
570 | 576 | kwargs.pop("kind", None) # e.g. "mergesort" is irrelevant
|
571 | 577 | nv.validate_argsort(args, kwargs)
|
572 | 578 |
|
| 579 | + start, stop, step = None, None, None |
573 | 580 | if self._range.step > 0:
|
574 |
| - result = np.arange(len(self), dtype=np.intp) |
| 581 | + if ascending: |
| 582 | + start = len(self) |
| 583 | + else: |
| 584 | + start, stop, step = len(self) - 1, -1, -1 |
| 585 | + elif ascending: |
| 586 | + start, stop, step = len(self) - 1, -1, -1 |
575 | 587 | else:
|
576 |
| - result = np.arange(len(self) - 1, -1, -1, dtype=np.intp) |
| 588 | + start = len(self) |
577 | 589 |
|
578 |
| - if not ascending: |
579 |
| - result = result[::-1] |
580 |
| - return result |
| 590 | + return np.arange(start, stop, step, dtype=np.intp) |
581 | 591 |
|
582 | 592 | def factorize(
|
583 | 593 | self,
|
584 | 594 | sort: bool = False,
|
585 | 595 | use_na_sentinel: bool = True,
|
586 | 596 | ) -> tuple[npt.NDArray[np.intp], RangeIndex]:
|
587 |
| - codes = np.arange(len(self), dtype=np.intp) |
588 |
| - uniques = self |
589 | 597 | if sort and self.step < 0:
|
590 |
| - codes = codes[::-1] |
591 |
| - uniques = uniques[::-1] |
| 598 | + codes = np.arange(len(self) - 1, -1, -1, dtype=np.intp) |
| 599 | + uniques = self[::-1] |
| 600 | + else: |
| 601 | + codes = np.arange(len(self), dtype=np.intp) |
| 602 | + uniques = self |
592 | 603 | return codes, uniques
|
593 | 604 |
|
594 | 605 | def equals(self, other: object) -> bool:
|
@@ -699,26 +710,15 @@ def _intersection(self, other: Index, sort: bool = False):
|
699 | 710 | # intersection disregarding the lower bounds
|
700 | 711 | tmp_start = first.start + (second.start - first.start) * first.step // gcd * s
|
701 | 712 | new_step = first.step * second.step // gcd
|
702 |
| - new_range = range(tmp_start, int_high, new_step) |
703 |
| - new_index = self._simple_new(new_range) |
704 | 713 |
|
705 | 714 | # adjust index to limiting interval
|
706 |
| - new_start = new_index._min_fitting_element(int_low) |
707 |
| - new_range = range(new_start, new_index.stop, new_index.step) |
708 |
| - new_index = self._simple_new(new_range) |
| 715 | + new_start = min_fitting_element(tmp_start, new_step, int_low) |
| 716 | + new_range = range(new_start, int_high, new_step) |
709 | 717 |
|
710 |
| - if (self.step < 0 and other.step < 0) is not (new_index.step < 0): |
711 |
| - new_index = new_index[::-1] |
| 718 | + if (self.step < 0 and other.step < 0) is not (new_range.step < 0): |
| 719 | + new_range = new_range[::-1] |
712 | 720 |
|
713 |
| - if sort is None: |
714 |
| - new_index = new_index.sort_values() |
715 |
| - |
716 |
| - return new_index |
717 |
| - |
718 |
| - def _min_fitting_element(self, lower_limit: int) -> int: |
719 |
| - """Returns the smallest element greater than or equal to the limit""" |
720 |
| - no_steps = -(-(lower_limit - self.start) // abs(self.step)) |
721 |
| - return self.start + abs(self.step) * no_steps |
| 721 | + return self._simple_new(new_range) |
722 | 722 |
|
723 | 723 | def _extended_gcd(self, a: int, b: int) -> tuple[int, int, int]:
|
724 | 724 | """
|
@@ -904,9 +904,9 @@ def _difference(self, other, sort=None):
|
904 | 904 | # e.g. range(10) and range(0, 10, 3)
|
905 | 905 | return super()._difference(other, sort=sort)
|
906 | 906 |
|
907 |
| - new_index = type(self)._simple_new(new_rng, name=res_name) |
908 | 907 | if first is not self._range:
|
909 |
| - new_index = new_index[::-1] |
| 908 | + new_rng = new_rng[::-1] |
| 909 | + new_index = type(self)._simple_new(new_rng, name=res_name) |
910 | 910 |
|
911 | 911 | return new_index
|
912 | 912 |
|
|
0 commit comments