#17 median_of_3_partition

Author: wojtask

src/book/chapter7/section4.py

@@ -1,6 +1,7 @@
 from book.chapter7.section1 import partition
 from book.data_structures import Array
 from book.data_structures import CT
+from solutions.chapter5.section1.exercise2 import random
 from util import range_of
 
 
@@ -26,3 +27,28 @@ def insertion_quicksort(A: Array[CT], p: int, r: int, k: int) -> None:
                     A[j + 1] = A[j]
                     j -= 1
                 A[j + 1] = key
+
+
+def median_of_3_partition(A: Array[CT], p: int, r: int) -> int:
+    """Partitions an array into two subarrays, the low side and the high side, such that each element in the low side of
+    the partition is less than or equal to the pivot value, which is, in turn, less than or equal to each element in the
+    high side. Uses a median of randomly picked three elements as the pivot.
+
+    Args:
+        A: an Array to partition
+        p: the lower index of the subarray to partition
+        r: the upper index of the subarray to partition
+
+    Returns:
+        The index q, such that each element in A[p:q - 1] is less than or equal to A[q], and that A[q] is less than or
+        equal to each element in A[q + 1:r].
+    """
+    i1, i2, i3 = random(p, r), random(p, r), random(p, r)
+    if A[i2] <= A[i1] <= A[i3] or A[i3] <= A[i1] <= A[i2]:
+        m = i1
+    elif A[i1] <= A[i2] <= A[i3] or A[i3] <= A[i2] <= A[i1]:
+        m = i2
+    else:
+        m = i3
+    A[m], A[r] = A[r], A[m]
+    return partition(A, p, r)

test/test_book/test_chapter7.py

@@ -6,8 +6,10 @@
 from book.chapter7.section1 import quicksort
 from book.chapter7.section3 import randomized_quicksort
 from book.chapter7.section4 import insertion_quicksort
+from book.chapter7.section4 import median_of_3_partition
 from test_case import ClrsTestCase
 from test_util import create_array
+from util import range_of
 
 
 class TestChapter7(ClrsTestCase):
@@ -45,3 +47,17 @@ def test_insertion_quicksort(self, data):
 
         self.assertArraySorted(A, end=n)
         self.assertArrayPermuted(A, elements, end=n)
+
+    @given(st.data())
+    def test_median_of_3_partition(self, data):
+        elements = data.draw(lists(integers(), min_size=1))
+        A = create_array(elements)
+        n = len(elements)
+
+        actual_pivot_index = median_of_3_partition(A, 1, n)
+
+        for i in range_of(1, to=actual_pivot_index):
+            self.assertLessEqual(A[i], A[actual_pivot_index])
+        for i in range_of(actual_pivot_index + 1, to=n):
+            self.assertGreaterEqual(A[i], A[actual_pivot_index])
+        self.assertArrayPermuted(A, elements, end=n)