LeetCode 4. Median of Two Sorted Arrays

Post by ailswan Aug. 31, 2024

4. Median of Two Sorted Arrays

Problem Statement

link: LeetCode.cn LeetCode 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:

Input: nums1 = [1,3], nums2 = [2] Output: 2.00000

Input: nums1 = [1,2], nums2 = [3,4] Output: 2.50000

Constraints: nums1.length == m nums2.length == n 0 <= m <= 1000 0 <= n <= 1000 1 <= m + n <= 2000 -106 <= nums1[i], nums2[i] <= 10^6

Solution Approach

Algorithm

Complexity

time complexity: O(log(min(m, n))),where m and n are the lengths of the two input arrays. space complexity: O(1).

Implement

    class Solution:
        def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
            def getKthElement(k):
                index1, index2 = 0, 0
                while True:
                    if index1 == m:
                        return nums2[index2 + k - 1]
                    if index2 == n:
                        return nums1[index1 + k - 1]
                    if k == 1:
                        return min(nums1[index1], nums2[index2])
                        
                    newIndex1 = min(index1 + k // 2 - 1, m - 1)
                    newIndex2 = min(index2 + k // 2 - 1, n - 1)
                    p1, p2 = nums1[newIndex1], nums2[newIndex2]
                    if p1 < p2:
                        k -= (newIndex1 - index1 + 1)
                        index1 = newIndex1 + 1
                    else:
                        k -= newIndex2 - index2 + 1
                        index2 = newIndex2 + 1
            
            m, n = len(nums1), len(nums2)
            totalLength = m + n
            if totalLength % 2 == 1:
                return getKthElement((totalLength + 1) // 2)
            else:
                return (getKthElement(totalLength // 2) + getKthElement(totalLength // 2 + 1)) / 2