0004: Median of Two Sorted Arrays

Problem Statement

Given two sorted arrays nums1 and nums2 of size m and n respectively, return the median of the two sorted arrays.

The overall run time complexity should be O(log (m+n)).

Example 1:

Input: nums1 = [1,3], nums2 = [2]
Output: 2.00000
Explanation: merged array = [1,2,3] and median is 2.

Example 2:

Input: nums1 = [1,2], nums2 = [3,4]
Output: 2.50000
Explanation: merged array = [1,2,3,4] and median is (2 + 3) / 2 = 2.5.

Constraints:

  • nums1.length == m
  • nums2.length == n
  • 0 <= m <= 1000
  • 0 <= n <= 1000
  • 1 <= m + n <= 2000
  • -106 <= nums1[i], nums2[i] <= 106

Code Solution

class Solution:
	
    def get_extremes(self, arr, p, n):
        lmax = float('-inf') if p == 0 else arr[p - 1]
        rmin = float('+inf') if p == n else arr[p]
        return lmax, rmin
		
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        # s1: assure len(nums1) > len(nums2)
        # s2: figure out pivots
        # s3: find boundary elements (`upr_lmax`, `upr_rmin`, `btm_lmax`, `btm_rmin`)
        # s4: converge and return median
        
        
        # s1: assure len(nums1) > len(nums2)
        if len(nums1) > len(nums2):
            nums1, nums2 = nums2, nums1
			
        nU, nL = len(nums1), len(nums2)
        
        lx = 0
        rx = nU
        while lx <= rx:
	        
            # s2: calculate pivots:
            p1 = (lx + rx) // 2
            p2 = (nU + nL + 1) // 2 - p1
            
            # s3: find boundaries:
            upr_lmax, upr_rmin = self.get_extremes(nums1, p1, nU)
            btm_lmax, btm_rmin = self.get_extremes(nums2, p2, nL)            
            
            # s4: converge and return median:
            if max(upr_lmax, btm_lmax) <= min(upr_rmin, btm_rmin):
                if (nL + nU) % 2 == 0:
                    return (max(upr_lmax, btm_lmax) + min(upr_rmin, btm_rmin)) / 2
                else:
                    return max(upr_lmax, btm_lmax)
            else:
                if btm_rmin < upr_lmax:
                    rx = p1 - 1
                elif upr_rmin < btm_lmax:
                    lx = p1 + 1		    
        return -1