Median of Two Sorted Arrays in Python

class Solution:
    # A function that takes two sorted arrays of numbers and returns the median of the two arrays
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        # Get the lengths of the two arrays
        m = len(nums1)
        n = len(nums2)

        # If the first array is longer than the second, swap them
        if m > n:
            nums1, nums2 = nums2, nums1
            m, n = n, m

        # Initialize the binary search variables
        i_min = 0  # The minimum index of the first array
        i_max = m  # The maximum index of the first array
        half_len = (m + n + 1) // 2  # The half length of the merged array

        # Perform binary search to find the partition point of the two arrays
        while i_min <= i_max:
            # The partition index of the first array
            i = (i_min + i_max) // 2
            # The partition index of the second array
            j = half_len - i

            # If the left element of the first array is too small, move the partition to the right
            if i < i_max and nums2[j - 1] > nums1[i]:
                i_min = i + 1
            # If the left element of the first array is too large, move the partition to the left
            elif i > i_min and nums1[i - 1] > nums2[j]:
                i_max = i - 1
            # Otherwise, we have found the correct partition
            else:
                # The maximum left element of the merged array
                max_left = 0
                # If the first array has no left element, use the second array's left element
                if i == 0:
                    max_left = nums2[j - 1]
                # If the second array has no left element, use the first array's left element
                elif j == 0:
                    max_left = nums1[i - 1]
                # Otherwise, use the larger of the two left elements
                else:
                    max_left = max(nums1[i - 1], nums2[j - 1])

                # If the total length of the merged array is odd, return the maximum left element as the median
                if (m + n) % 2 == 1:
                    return max_left

                # The minimum right element of the merged array
                min_right = 0
                # If the first array has no right element, use the second array's right element
                if i == m:
                    min_right = nums2[j]
                # If the second array has no right element, use the first array's right element
                elif j == n:
                    min_right = nums1[i]
                # Otherwise, use the smaller of the two right elements
                else:
                    min_right = min(nums1[i], nums2[j])

                # If the total length of the merged array is even, return the average of the maximum left and minimum right elements as the median
                return (max_left + min_right) / 2

        # If the binary search fails, return 0 as a default value
        return 0

PrevNext