diff --git a/maths/median_of_medians.py b/maths/median_of_medians.py new file mode 100644 index 000000000000..734c47d953fd --- /dev/null +++ b/maths/median_of_medians.py @@ -0,0 +1,69 @@ +""" +Median of Medians Algorithm +Guarantees O(n) worst-case time complexity for finding the k-th smallest element. +Reference: https://en.wikipedia.org/wiki/Median_of_medians +""" + + +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] + 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, rank: int) -> int: + """ + Find the k-th smallest element in an unsorted list using Median of Medians. + + Args: + arr: List of comparable elements + rank: 1-based index of the desired smallest element + + Returns: + The k-th smallest element in arr + + Raises: + 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 <= rank <= len(arr): + msg = f"rank={rank} is out of range for array of length {len(arr)}" + raise ValueError(msg) + + if len(arr) <= 5: + return sorted(arr)[rank - 1] + + chunks = [arr[i : i + 5] for i in range(0, len(arr), 5)] + medians = [sorted(chunk)[len(chunk) // 2] for chunk in chunks] + + pivot = median_of_medians(medians, len(medians) // 2 + 1) + + low, equal, high = partition(arr, pivot) + + if rank <= len(low): + return median_of_medians(low, rank) + elif rank <= len(low) + len(equal): + return pivot + else: + return median_of_medians(high, rank - len(low) - len(equal)) + + +if __name__ == "__main__": + import doctest + + doctest.testmod()