1+ """
2+ Median of Medians Algorithm
3+ Guarantees O(n) worst-case time complexity for finding the k-th smallest element.
4+ Reference: https://en.wikipedia.org/wiki/Median_of_medians
5+ """
6+
7+
8+ def partition (arr : list , pivot : int ) -> tuple :
9+ """Partition array into elements less than, equal to, and greater than pivot."""
10+ low = [x for x in arr if x < pivot ]
11+ high = [x for x in arr if x > pivot ]
12+ equal = [x for x in arr if x == pivot ]
13+ return low , equal , high
14+
15+
16+ def median_of_medians (arr : list , k : int ) -> int :
17+ """
18+ Find the k-th smallest element in an unsorted list using Median of Medians.
19+
20+ Args:
21+ arr: List of comparable elements
22+ k: 1-based index of the desired smallest element
23+
24+ Returns:
25+ The k-th smallest element in arr
26+
27+ Raises:
28+ ValueError: If k is out of range
29+
30+ Examples:
31+ >>> median_of_medians([3, 1, 4, 1, 5, 9, 2, 6], 3)
32+ 2
33+ >>> median_of_medians([7, 2, 10, 5], 1)
34+ 2
35+ >>> median_of_medians([1, 2, 3, 4, 5], 5)
36+ 5
37+ """
38+ if not 1 <= k <= len (arr ):
39+ raise ValueError (f"k={ k } is out of range for array of length { len (arr )} " )
40+
41+ # Base case
42+ if len (arr ) <= 5 :
43+ return sorted (arr )[k - 1 ]
44+
45+ # Step 1: Divide into chunks of 5 and find median of each chunk
46+ chunks = [arr [i : i + 5 ] for i in range (0 , len (arr ), 5 )]
47+ medians = [sorted (chunk )[len (chunk ) // 2 ] for chunk in chunks ]
48+
49+ # Step 2: Recursively find median of medians
50+ pivot = median_of_medians (medians , len (medians ) // 2 + 1 )
51+
52+ # Step 3: Partition around pivot
53+ low , equal , high = partition (arr , pivot )
54+
55+ # Step 4: Recurse into the correct partition
56+ if k <= len (low ):
57+ return median_of_medians (low , k )
58+ elif k <= len (low ) + len (equal ):
59+ return pivot
60+ else :
61+ return median_of_medians (high , k - len (low ) - len (equal ))
62+
63+
64+ if __name__ == "__main__" :
65+ import doctest
66+ doctest .testmod ()
67+
68+ sample = [3 , 1 , 4 , 1 , 5 , 9 , 2 , 6 , 5 , 3 , 5 ]
69+ print (f"Array: { sample } " )
70+ for i in range (1 , len (sample ) + 1 ):
71+ print (f" { i } -th smallest: { median_of_medians (sample , i )} " )
0 commit comments