diff --git a/problems/0004-median-of-two-sorted-arrays/solution.py b/problems/0004-median-of-two-sorted-arrays/solution.py new file mode 100644 index 0000000..91156a3 --- /dev/null +++ b/problems/0004-median-of-two-sorted-arrays/solution.py @@ -0,0 +1,24 @@ +class Solution: + def findMedianSortedArrays(self, nums1: list[int], nums2: list[int]) -> float: + if len(nums1) <= len(nums2): + A, B = nums1, nums2 + else: + A, B = nums2, nums1 + half_total = (len(A) + len(B)) // 2 + lo, hi = 0, len(A) + while True: + i = (lo + hi) // 2 + j = half_total - i + Aleft = A[i - 1] if i > 0 else float("-inf") + Aright = A[i] if i < len(A) else float("inf") + Bleft = B[j - 1] if j > 0 else float("-inf") + Bright = B[j] if j < len(B) else float("inf") + if Aleft > Bright: + hi = i - 1 + elif Bleft > Aright: + lo = i + 1 + else: + if (len(A) + len(B)) % 2 != 0: + return min(Aright, Bright) + else: + return (max(Aleft, Bleft) + min(Aright, Bright)) / 2 diff --git a/problems/0004-median-of-two-sorted-arrays/solution_test.py b/problems/0004-median-of-two-sorted-arrays/solution_test.py new file mode 100644 index 0000000..6e00963 --- /dev/null +++ b/problems/0004-median-of-two-sorted-arrays/solution_test.py @@ -0,0 +1,13 @@ +import pytest +from solution import Solution + + +@pytest.mark.parametrize( + "nums1, nums2, expected", + [ + ([1, 3], [2], 2), + ([1, 2], [3, 4], 2.5), + ], +) +def test_REPLACE(nums1, nums2, expected): + assert Solution().findMedianSortedArrays(nums1, nums2) == expected