|
19 | 19 | ensure_float64, ensure_int64, ensure_object, ensure_platform_int, |
20 | 20 | ensure_uint64, is_array_like, is_bool_dtype, is_categorical_dtype, |
21 | 21 | is_complex_dtype, is_datetime64_any_dtype, is_datetime64tz_dtype, |
22 | | - is_datetimelike, is_extension_array_dtype, is_float_dtype, |
| 22 | + is_datetimelike, is_extension_array_dtype, is_float_dtype, is_integer, |
23 | 23 | is_integer_dtype, is_interval_dtype, is_list_like, is_numeric_dtype, |
24 | 24 | is_object_dtype, is_period_dtype, is_scalar, is_signed_integer_dtype, |
25 | 25 | is_sparse, is_timedelta64_dtype, is_unsigned_integer_dtype, |
@@ -1729,6 +1729,89 @@ def func(arr, indexer, out, fill_value=np.nan): |
1729 | 1729 | return out |
1730 | 1730 |
|
1731 | 1731 |
|
| 1732 | +# ------------ # |
| 1733 | +# searchsorted # |
| 1734 | +# ------------ # |
| 1735 | + |
| 1736 | +def searchsorted(arr, value, side="left", sorter=None): |
| 1737 | + """ |
| 1738 | + Find indices where elements should be inserted to maintain order. |
| 1739 | +
|
| 1740 | + .. versionadded:: 0.25.0 |
| 1741 | +
|
| 1742 | + Find the indices into a sorted array `arr` (a) such that, if the |
| 1743 | + corresponding elements in `value` were inserted before the indices, |
| 1744 | + the order of `arr` would be preserved. |
| 1745 | +
|
| 1746 | + Assuming that `arr` is sorted: |
| 1747 | +
|
| 1748 | + ====== ================================ |
| 1749 | + `side` returned index `i` satisfies |
| 1750 | + ====== ================================ |
| 1751 | + left ``arr[i-1] < value <= self[i]`` |
| 1752 | + right ``arr[i-1] <= value < self[i]`` |
| 1753 | + ====== ================================ |
| 1754 | +
|
| 1755 | + Parameters |
| 1756 | + ---------- |
| 1757 | + arr: array-like |
| 1758 | + Input array. If `sorter` is None, then it must be sorted in |
| 1759 | + ascending order, otherwise `sorter` must be an array of indices |
| 1760 | + that sort it. |
| 1761 | + value : array_like |
| 1762 | + Values to insert into `arr`. |
| 1763 | + side : {'left', 'right'}, optional |
| 1764 | + If 'left', the index of the first suitable location found is given. |
| 1765 | + If 'right', return the last such index. If there is no suitable |
| 1766 | + index, return either 0 or N (where N is the length of `self`). |
| 1767 | + sorter : 1-D array_like, optional |
| 1768 | + Optional array of integer indices that sort array a into ascending |
| 1769 | + order. They are typically the result of argsort. |
| 1770 | +
|
| 1771 | + Returns |
| 1772 | + ------- |
| 1773 | + array of ints |
| 1774 | + Array of insertion points with the same shape as `value`. |
| 1775 | +
|
| 1776 | + See Also |
| 1777 | + -------- |
| 1778 | + numpy.searchsorted : Similar method from NumPy. |
| 1779 | + """ |
| 1780 | + if sorter is not None: |
| 1781 | + sorter = ensure_platform_int(sorter) |
| 1782 | + |
| 1783 | + if isinstance(arr, np.ndarray) and is_integer_dtype(arr) and ( |
| 1784 | + is_integer(value) or is_integer_dtype(value)): |
| 1785 | + from .arrays.array_ import array |
| 1786 | + # if `arr` and `value` have different dtypes, `arr` would be |
| 1787 | + # recast by numpy, causing a slow search. |
| 1788 | + # Before searching below, we therefore try to give `value` the |
| 1789 | + # same dtype as `arr`, while guarding against integer overflows. |
| 1790 | + iinfo = np.iinfo(arr.dtype.type) |
| 1791 | + value_arr = np.array([value]) if is_scalar(value) else np.array(value) |
| 1792 | + if (value_arr >= iinfo.min).all() and (value_arr <= iinfo.max).all(): |
| 1793 | + # value within bounds, so no overflow, so can convert value dtype |
| 1794 | + # to dtype of arr |
| 1795 | + dtype = arr.dtype |
| 1796 | + else: |
| 1797 | + dtype = value_arr.dtype |
| 1798 | + |
| 1799 | + if is_scalar(value): |
| 1800 | + value = dtype.type(value) |
| 1801 | + else: |
| 1802 | + value = array(value, dtype=dtype) |
| 1803 | + elif not (is_object_dtype(arr) or is_numeric_dtype(arr) or |
| 1804 | + is_categorical_dtype(arr)): |
| 1805 | + from pandas.core.series import Series |
| 1806 | + # E.g. if `arr` is an array with dtype='datetime64[ns]' |
| 1807 | + # and `value` is a pd.Timestamp, we may need to convert value |
| 1808 | + value_ser = Series(value)._values |
| 1809 | + value = value_ser[0] if is_scalar(value) else value_ser |
| 1810 | + |
| 1811 | + result = arr.searchsorted(value, side=side, sorter=sorter) |
| 1812 | + return result |
| 1813 | + |
| 1814 | + |
1732 | 1815 | # ---- # |
1733 | 1816 | # diff # |
1734 | 1817 | # ---- # |
|
0 commit comments