|
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