Skip to content

Commit

Permalink
Add QuickSelect algorithm (yangshun#66)
Browse files Browse the repository at this point in the history
* QuickSelect implementation

* QuickSelect tests
  • Loading branch information
louietyj authored and yangshun committed Oct 20, 2017
1 parent 856b1e6 commit 4c672de
Showing 1 changed file with 60 additions and 0 deletions.
60 changes: 60 additions & 0 deletions utilities/python/quick_select.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 4c672de

Please sign in to comment.