Python/searches/quick_select.py

63 lines
1.7 KiB
Python
Raw Normal View History

2017-10-09 21:26:27 +00:00
"""
A Python implementation of the quick select algorithm, which is efficient for
calculating the value that would appear in the index of a list if it would be
sorted, even if it is not already sorted
2017-10-09 21:26:27 +00:00
https://en.wikipedia.org/wiki/Quickselect
"""
import random
2019-10-05 05:14:13 +00:00
def _partition(data: list, pivot) -> tuple:
2017-10-09 21:26:27 +00:00
"""
Three way partition the data into smaller, equal and greater lists,
in relationship to the pivot
:param data: The data to be sorted (a list)
:param pivot: The value to partition the data on
:return: Three list: smaller, equal and greater
"""
less, equal, greater = [], [], []
for element in data:
if element < pivot:
2017-10-09 21:26:27 +00:00
less.append(element)
elif element > pivot:
2017-10-09 21:26:27 +00:00
greater.append(element)
else:
equal.append(element)
return less, equal, greater
2019-10-05 05:14:13 +00:00
def quick_select(items: list, index: int):
"""
>>> quick_select([2, 4, 5, 7, 899, 54, 32], 5)
54
>>> quick_select([2, 4, 5, 7, 899, 54, 32], 1)
4
>>> quick_select([5, 4, 3, 2], 2)
4
>>> quick_select([3, 5, 7, 10, 2, 12], 3)
7
"""
# index = len(items) // 2 when trying to find the median
# (value of index when items is sorted)
2019-10-05 05:14:13 +00:00
# invalid input
if index >= len(items) or index < 0:
return None
2019-10-05 05:14:13 +00:00
pivot = random.randint(0, len(items) - 1)
pivot = items[pivot]
count = 0
smaller, equal, larger = _partition(items, pivot)
count = len(equal)
m = len(smaller)
2017-10-09 21:26:27 +00:00
# index is the pivot
if m <= index < m + count:
2017-10-09 21:26:27 +00:00
return pivot
# must be in smaller
elif m > index:
return quick_select(smaller, index)
2019-10-05 05:14:13 +00:00
# must be in larger
else:
return quick_select(larger, index - (m + count))