Median of two sorted arrays

median_of_sorted_arrays.py

from typing import List
from random import randint


def findMedianSortedArrays(nums1: List[int], nums2: List[int]) -> float:
    n1: int = len(nums1)
    n2: int = len(nums2)

    if n1 > n2:
        return findMedianSortedArrays(nums2, nums1)

    total: int = n1 + n2
    k: int = (total + 1) // 2
    odd: bool = bool(total & 1)
    l: int = 0
    r: int = n1
    while l < r:
        m1: int = l + (r - l) // 2
        m2: int = k - m1
        if nums1[m1] < nums2[m2 - 1]:
            l = m1 + 1
        else:
            r = m1

    pairs = ((nums1, l), (nums2, k - l))
    c1 = max(arr[x - 1] for arr, x in pairs if x > 0)

    if odd:
        return c1
    else:
        c2 = min(arr[x] for arr, x in pairs if x < len(arr))
        return (c1 + c2) * 0.5


def median(arr: List[int]) -> float:
    x: int = len(arr)
    y: int = x // 2
    if x & 1:
        return arr[y]
    else:
        return (arr[y - 1] + arr[y]) / 2.0


def main():
    for length in range(1, 50):
        for _ in range(1000):
            r1 = sorted(randint(0, 100) for _ in range(length))
            r2 = sorted(randint(0, 100) for _ in range(4))
            x = findMedianSortedArrays(r1, r2)
            y = median(sorted(r1 + r2))
            if x != y:
                print(r1)
                print(r2)
                print(x, y)
                return


if __name__ == '__main__':
    main()