beam: median of two sorted arrays (binary partition) + fuzz tests

wizard-work/beam/median_two_sorted.py

@@ -0,0 +1,76 @@
+"""Median of two sorted arrays (LeetCode 4).
+
+Returns the median of the two sorted input arrays in O(log(min(m,n))) time
+using the binary-search partition method.
+
+Edge cases handled:
+  - either array empty
+  - arrays of wildly different lengths
+  - duplicate values across arrays
+  - total length odd or even
+  - all elements of one array smaller than all of the other
+"""
+
+from __future__ import annotations
+
+
+def find_median_sorted_arrays(a: list[int], b: list[int]) -> float:
+    # Ensure `a` is the shorter array so we binary-search the smaller one.
+    if len(a) > len(b):
+        a, b = b, a
+
+    m, n = len(a), len(b)
+    # We want (m + n + 1) // 2 elements on the left side.
+    half = (m + n + 1) // 2
+
+    lo, hi = 0, m  # partition index in `a`: how many of `a` go left
+    while lo <= hi:
+        i = (lo + hi) // 2          # number taken from `a` on the left
+        j = half - i                # number taken from `b` on the left
+
+        # Edge sentinels: -inf / +inf when a partition is empty or full.
+        a_left = a[i - 1] if i > 0 else float("-inf")
+        a_right = a[i] if i < m else float("inf")
+        b_left = b[j - 1] if j > 0 else float("-inf")
+        b_right = b[j] if j < n else float("inf")
+
+        if a_left <= b_right and b_left <= a_right:
+            # Correct partition found.
+            if (m + n) % 2 == 1:
+                return float(max(a_left, b_left))
+            return (max(a_left, b_left) + min(a_right, b_right)) / 2.0
+        if a_left > b_right:
+            # Took too many from `a`; move left.
+            hi = i - 1
+        else:
+            # Took too few from `a`; move right.
+            lo = i + 1
+
+    raise ValueError("Inputs must be sorted ascending; no valid partition found")
+
+
+def _brute(a, b):
+    merged = sorted(a + b)
+    L = len(merged)
+    if L == 0:
+        raise ValueError("both empty")
+    if L % 2 == 1:
+        return float(merged[L // 2])
+    return (merged[L // 2 - 1] + merged[L // 2]) / 2.0
+
+
+if __name__ == "__main__":
+    # Self-check against brute force over a range of cases.
+    import random
+    rng = random.Random(42)
+    for _ in range(2000):
+        la = rng.randint(0, 12)
+        lb = rng.randint(0, 12)
+        a = sorted(rng.randint(-20, 20) for _ in range(la))
+        b = sorted(rng.randint(-20, 20) for _ in range(lb))
+        if la == 0 and lb == 0:
+            continue
+        got = find_median_sorted_arrays(a, b)
+        exp = _brute(a, b)
+        assert got == exp, (a, b, got, exp)
+    print("self-check: 2000 random cases passed")

wizard-work/beam/test_median.py

@@ -0,0 +1,31 @@
+import random
+import sys
+sys.path.insert(0, "/root/.hermes/wizard-outputs/beam")
+from median_two_sorted import find_median_sorted_arrays as f
+
+cases = [
+    ([1,3],[2],2.0),
+    ([1,2],[3,4],2.5),
+    ([],[1],1.0),
+    ([2],[],2.0),
+    ([1,2,3,4,5],[6,7,8,9,10],5.5),
+    ([0,0],[0,0],0.0),
+    ([-5,-3,-1],[ -2,0,2],-1.5),
+]
+for a,b,exp in cases:
+    got = f(a,b)
+    assert got == exp, (a,b,got,exp)
+
+# brute-force fuzz
+rng = random.Random(7)
+for _ in range(5000):
+    la = rng.randint(0, 15)
+    lb = rng.randint(0, 15)
+    a = sorted(rng.randint(-50,50) for _ in range(la))
+    b = sorted(rng.randint(-50,50) for _ in range(lb))
+    if la==0 and lb==0: continue
+    merged = sorted(a+b); L=len(merged)
+    exp = float(merged[L//2]) if L%2 else (merged[L//2-1]+merged[L//2])/2.0
+    got = f(a,b)
+    assert got == exp, (a,b,got,exp)
+print("ALL TESTS PASSED")