1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55
|
"""
Utility functions for sorting and dealing with sorted Series and DataFrames
"""
from bisect import bisect_left, bisect_right
from typing import Any, Callable
def sorted_exists(values: list, x: Any) -> tuple[bool, int]:
"""
For list, values, returns the insert position for item x and whether the item already exists in the list. This
allows one function call to return either the index to overwrite an existing value in the list, or the index to
insert a new item in the list and keep the list in sorted order.
:param values: list
:param x: item
:return: (exists, index) tuple
"""
i = bisect_left(values, x)
j = bisect_right(values, x)
exists = x in values[i:j]
return exists, i
def sorted_index(values: list, x: Any) -> int:
"""
For list, values, returns the index location of element x. If x does not exist will raise an error.
:param values: list
:param x: item
:return: integer index
"""
i = bisect_left(values, x)
j = bisect_right(values, x)
return values[i:j].index(x) + i
def sorted_list_indexes(list_to_sort: list, key: Callable | Any = None, reverse: bool = False) -> list[int]:
"""
Sorts a list but returns the order of the index values of the list for the sort and not the values themselves.
For example is the list provided is ['b', 'a', 'c'] then the result will be [2, 1, 3]
:param list_to_sort: list to sort
:param key: if not None then a function of one argument that is used to extract a comparison key from each
list element
:param reverse: if True then the list elements are sorted as if each comparison were reversed.
:return: list of sorted index values
"""
if key is not None:
def key_func(i):
return key(list_to_sort.__getitem__(i))
else:
key_func = list_to_sort.__getitem__
return sorted(range(len(list_to_sort)), key=key_func, reverse=reverse)
|