Update quick_select.py (#1523)

* Update quick_select.py

Add Doctests.

* Add typehints

* Don't pre-allocate "smaller" and "larger"
This commit is contained in:
percy07 2019-10-30 20:40:30 +05:30 committed by Christian Clauss
parent fc533a7598
commit df95f43907

View File

@ -1,12 +1,13 @@
import random
""" """
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 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
https://en.wikipedia.org/wiki/Quickselect https://en.wikipedia.org/wiki/Quickselect
""" """
import random
def _partition(data, pivot): def _partition(data: list, pivot) -> tuple:
""" """
Three way partition the data into smaller, equal and greater lists, Three way partition the data into smaller, equal and greater lists,
in relationship to the pivot in relationship to the pivot
@ -25,28 +26,37 @@ def _partition(data, pivot):
return less, equal, greater return less, equal, greater
def quickSelect(list, k): def quick_select(items: list, index: int):
# k = len(list) // 2 when trying to find the median (index that value would be when list is sorted) """
>>> 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)
# invalid input # invalid input
if k >= len(list) or k < 0: if index >= len(items) or index < 0:
return None return None
smaller = [] pivot = random.randint(0, len(items) - 1)
larger = [] pivot = items[pivot]
pivot = random.randint(0, len(list) - 1)
pivot = list[pivot]
count = 0 count = 0
smaller, equal, larger = _partition(list, pivot) smaller, equal, larger = _partition(items, pivot)
count = len(equal) count = len(equal)
m = len(smaller) m = len(smaller)
# k is the pivot # index is the pivot
if m <= k < m + count: if m <= index < m + count:
return pivot return pivot
# must be in smaller # must be in smaller
elif m > k: elif m > index:
return quickSelect(smaller, k) return quick_select(smaller, index)
# must be in larger # must be in larger
else: else:
return quickSelect(larger, k - (m + count)) return quick_select(larger, index - (m + count))