From 4c672de8e6b32925b692670898d81efcd20bff48 Mon Sep 17 00:00:00 2001
From: louietyj <louietyj@gmail.com>
Date: Fri, 20 Oct 2017 17:29:28 +0800
Subject: [PATCH] Add QuickSelect algorithm (#66)

* QuickSelect implementation

* QuickSelect tests
---
 utilities/python/quick_select.py | 60 ++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)
 create mode 100644 utilities/python/quick_select.py

diff --git a/utilities/python/quick_select.py b/utilities/python/quick_select.py
new file mode 100644
index 00000000..c121d0d4
--- /dev/null
+++ b/utilities/python/quick_select.py
@@ -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)