Files
troupe-standup/wizard-work/beam/test_median.py
T

32 lines
863 B
Python

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")