fix: add doctest to partition and rename k to rank

Author: shubhangiisinghh

maths/median_of_medians.py

@@ -5,68 +5,63 @@
 """
 
 
-def partition(arr: list, pivot: int) -> tuple:
-    """Partition array into elements less than, equal to, and greater than pivot."""
+def partition(arr: list, pivot: int) -> tuple[list, list, list]:
+    """
+    Partition array into elements less than, equal to, and greater than pivot.
+
+    >>> partition([3, 1, 4, 1, 5], 3)
+    ([1, 1], [3], [4, 5])
+    >>> partition([7, 7, 7], 7)
+    ([], [7, 7, 7], [])
+    """
     low = [x for x in arr if x < pivot]
-    high = [x for x in arr if x > pivot]
     equal = [x for x in arr if x == pivot]
+    high = [x for x in arr if x > pivot]
     return low, equal, high
 
 
-def median_of_medians(arr: list, k: int) -> int:
+def median_of_medians(arr: list, rank: int) -> int:
     """
     Find the k-th smallest element in an unsorted list using Median of Medians.
 
     Args:
         arr: List of comparable elements
-        k: 1-based index of the desired smallest element
+        rank: 1-based index of the desired smallest element
 
     Returns:
         The k-th smallest element in arr
 
     Raises:
-        ValueError: If k is out of range
-
-    Examples:
-        >>> median_of_medians([3, 1, 4, 1, 5, 9, 2, 6], 3)
-        2
-        >>> median_of_medians([7, 2, 10, 5], 1)
-        2
-        >>> median_of_medians([1, 2, 3, 4, 5], 5)
-        5
+        ValueError: If rank is out of range
+
+    >>> median_of_medians([3, 1, 4, 1, 5, 9, 2, 6], 3)
+    2
+    >>> median_of_medians([7, 2, 10, 5], 1)
+    2
+    >>> median_of_medians([1, 2, 3, 4, 5], 5)
+    5
     """
-    if not 1 <= k <= len(arr):
-        raise ValueError(f"k={k} is out of range for array of length {len(arr)}")
+    if not 1 <= rank <= len(arr):
+        raise ValueError(f"rank={rank} is out of range for array of length {len(arr)}")
 
-    # Base case
     if len(arr) <= 5:
-        return sorted(arr)[k - 1]
+        return sorted(arr)[rank - 1]
 
-    # Step 1: Divide into chunks of 5 and find median of each chunk
-    chunks = [arr[i : i + 5] for i in range(0, len(arr), 5)]
+    chunks = [arr[i: i + 5] for i in range(0, len(arr), 5)]
     medians = [sorted(chunk)[len(chunk) // 2] for chunk in chunks]
 
-    # Step 2: Recursively find median of medians
     pivot = median_of_medians(medians, len(medians) // 2 + 1)
 
-    # Step 3: Partition around pivot
     low, equal, high = partition(arr, pivot)
 
-    # Step 4: Recurse into the correct partition
-    if k <= len(low):
-        return median_of_medians(low, k)
-    elif k <= len(low) + len(equal):
+    if rank <= len(low):
+        return median_of_medians(low, rank)
+    elif rank <= len(low) + len(equal):
         return pivot
     else:
-        return median_of_medians(high, k - len(low) - len(equal))
+        return median_of_medians(high, rank - len(low) - len(equal))
 
 
 if __name__ == "__main__":
     import doctest
-
-    doctest.testmod()
-
-    sample = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
-    print(f"Array: {sample}")
-    for i in range(1, len(sample) + 1):
-        print(f"  {i}-th smallest: {median_of_medians(sample, i)}")
+    doctest.testmod()