forked from yangshun/front-end-interview-handbook
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add QuickSelect algorithm (yangshun#66)
* QuickSelect implementation * QuickSelect tests
- Loading branch information
Showing
1 changed file
with
60 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |