QuickSelect implementation

utilities/python/quick_select.py

@@ -0,0 +1,46 @@
+## 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