Skip to content

Commit df95f43

Browse files
percy07cclauss
authored andcommitted
Update quick_select.py (#1523)
* Update quick_select.py Add Doctests. * Add typehints * Don't pre-allocate "smaller" and "larger"
1 parent fc533a7 commit df95f43

File tree

1 file changed

+27
-17
lines changed

1 file changed

+27
-17
lines changed

searches/quick_select.py

+27-17
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
1-
import random
2-
31
"""
4-
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
2+
A Python implementation of the quick select algorithm, which is efficient for
3+
calculating the value that would appear in the index of a list if it would be
4+
sorted, even if it is not already sorted
55
https://en.wikipedia.org/wiki/Quickselect
66
"""
7+
import random
78

89

9-
def _partition(data, pivot):
10+
def _partition(data: list, pivot) -> tuple:
1011
"""
1112
Three way partition the data into smaller, equal and greater lists,
1213
in relationship to the pivot
@@ -25,28 +26,37 @@ def _partition(data, pivot):
2526
return less, equal, greater
2627

2728

28-
def quickSelect(list, k):
29-
# k = len(list) // 2 when trying to find the median (index that value would be when list is sorted)
29+
def quick_select(items: list, index: int):
30+
"""
31+
>>> quick_select([2, 4, 5, 7, 899, 54, 32], 5)
32+
54
33+
>>> quick_select([2, 4, 5, 7, 899, 54, 32], 1)
34+
4
35+
>>> quick_select([5, 4, 3, 2], 2)
36+
4
37+
>>> quick_select([3, 5, 7, 10, 2, 12], 3)
38+
7
39+
"""
40+
# index = len(items) // 2 when trying to find the median
41+
# (value of index when items is sorted)
3042

3143
# invalid input
32-
if k >= len(list) or k < 0:
44+
if index >= len(items) or index < 0:
3345
return None
3446

35-
smaller = []
36-
larger = []
37-
pivot = random.randint(0, len(list) - 1)
38-
pivot = list[pivot]
47+
pivot = random.randint(0, len(items) - 1)
48+
pivot = items[pivot]
3949
count = 0
40-
smaller, equal, larger = _partition(list, pivot)
50+
smaller, equal, larger = _partition(items, pivot)
4151
count = len(equal)
4252
m = len(smaller)
4353

44-
# k is the pivot
45-
if m <= k < m + count:
54+
# index is the pivot
55+
if m <= index < m + count:
4656
return pivot
4757
# must be in smaller
48-
elif m > k:
49-
return quickSelect(smaller, k)
58+
elif m > index:
59+
return quick_select(smaller, index)
5060
# must be in larger
5161
else:
52-
return quickSelect(larger, k - (m + count))
62+
return quick_select(larger, index - (m + count))

0 commit comments

Comments
 (0)