Skip to content

Commit e7c4505

Browse files
add: median of medians algorithm in maths
1 parent 841e947 commit e7c4505

File tree

1 file changed

+71
-0
lines changed

1 file changed

+71
-0
lines changed

maths/median_of_medians.py

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
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

Comments
 (0)