diff --git a/utilities/python/quick_select.py b/utilities/python/quick_select.py new file mode 100644 index 00000000..c121d0d4 --- /dev/null +++ b/utilities/python/quick_select.py @@ -0,0 +1,60 @@ +## QuickSelect -- Linear-time k-th order statistic +## (i.e. select the k-th smallest element in an unsorted array) +## https://en.wikipedia.org/wiki/Quickselect + +def partition(array, start, end, pivot): + """Partitions by a pivot value, which might not necessarily be in the array. + This variant is useful when you want to bound your recursion depth by the + range of the input values, and not the length of the array.""" + pivot_index = start + for i in range(start, end): + if array[i] <= pivot: + array[i], array[pivot_index] = array[pivot_index], array[i] + pivot_index += 1 + return pivot_index + +import random +def partition_first(array, start, end): + """Selects the first element as pivot. Returns the index where the pivot went to. + In this variant, we can guarantee that the pivot will be in its final sorted position. + We need this guarantee for QuickSelect.""" + if start + 1 == end: + return start + pivot = array[start] + pivot_index = start + 1 + for i in range(start + 1, end): + if array[i] <= pivot: + array[i], array[pivot_index] = array[pivot_index], array[i] + pivot_index += 1 + # Move pivot to front + array[start], array[pivot_index - 1] = array[pivot_index - 1], array[start] + return pivot_index - 1 + +def quick_select(array, k): + """NOTE: k-th smallest element counts from 0!""" + left = 0 + right = len(array) + while True: + random_index = random.sample(range(left, right), 1)[0] + array[left], array[random_index] = array[random_index], array[left] + pivot_index = partition_first(array, left, right) + if k == pivot_index: + return array[pivot_index] + if k < pivot_index: + right = pivot_index + else: + left = pivot_index + 1 + + + +print(quick_select([0], 0) == 0) +print(quick_select([0, 1, 2, 3, 4], 2) == 2) +print(quick_select([4, 3, 2, 1, 0], 2) == 2) +print(quick_select([1, 3, 4, 2, 0], 2) == 2) + +# Large test case, for randomized tests +lst = list(range(1000)) +for _ in range(10): + k = random.randint(0, 999) + random.shuffle(lst) + print(quick_select(lst, k) == k)