Add QuickSelect algorithm (#66)

* QuickSelect implementation

* QuickSelect tests
pull/68/head
louietyj 7 years ago committed by Yangshun Tay
parent 856b1e611e
commit 4c672de8e6

@ -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)
Loading…
Cancel
Save