Median of Two Sorted Arrays

Posted 2020-08-12 Sheryl Wang

There are two sorted arrays nums1 and nums2 of size m and n respectively.

Find the median of the two sorted arrays. The overall run time complexity should be O(log (m+n)).

Example 1:

nums1 = [1, 3]
nums2 = [2]

The median is 2.0

Example 2:

nums1 = [1, 2]
nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

虽然是hard,但是这题是必会高频题.主要思路为间接的二分.参考这个博客

理思路,假设A为m个,B为n个,如果m+n为奇数, 则直接查找A和B中的第(m+n)/2+1个元素,如果m+n为偶数,则先求第(m+n)/2和第(m+n)/2+1个元素,之后求平均->所以此题转化为Kth element in 2 sorted array.

而对于“Kth element in 2 sorted array”， 如下图,两个中位数 A[m/2] 和 B[n/2]， 可以将数组划分为四个部分。而丢弃哪一个部分取决于两个条件：1， (m/2 + n/2)和k的关系；2，A[m/2] 和B[n/2]的关系;

 

 

 

如果 (m/2 + n/2) > k，那么意味着，当前中位数取高了，正确的中位数要么在 Section 1或者Section3中。如果A[m/2] > B[n/2], 意味着中位数肯定不可能在Section 2里面，那么新的搜索可以丢弃这个区间段了。同理可以推断出余下三种情况，如下所示：

If (m/2+n/2+1) > k && am/2 > bn/2 , drop Section 2

If (m/2+n/2+1) > k && am/2 < bn/2 , drop Section 4

If (m/2+n/2+1) < k && am/2 > bn/2 ,  drop Section 3

If (m/2+n/2+1) < k && am/2 < bn/2 ,  drop Section 1

代码如下: 加粗部分一定要注意,注意只有加起来的数小于k时,我们才可以确定的将最小的那部分删除,同时缩小K.

class Solution(object):
    def findMedianSortedArrays(self, nums1, nums2):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float
        """
        m = len(nums1)
        n = len(nums2)
        if (m+n)&1:
            return self.findKthSortedArrays(nums1,0, m, nums2, 0, n, (m+n)/2+1)  
        else:
            return (self.findKthSortedArrays(nums1,0, m, nums2, 0, n, (m+n)/2) 
            + self.findKthSortedArrays(nums1,0, m, nums2, 0, n, (m+n)/2+1))/2.0
       
    def findKthSortedArrays(self, nums1, startA, lenA, nums2, startB, lenB, k):
        print startA, lenA, startB, lenB, k
        if lenA<=0:
            return nums2[startB+k-1]
        if lenB<=0:
            return nums1[startA+k-1]
        if k==1:
            return min(nums1[startA], nums2[startB])
        m1 = startA+lenA/2
        m2 = startB+lenB/2
        if nums1[m1] >= nums2[m2]:
            if lenA/2+lenB/2+1 >= k:
                return self.findKthSortedArrays(nums1, startA, lenA/2, nums2, startB, lenB, k)  #delete the latter part of A
            else:
                return self.findKthSortedArrays(nums1, startA, lenA, nums2, startB+lenB/2+1, lenB-lenB/2-1, k-lenB/2-1)  #delete the start part of B
        else:
            if m1 + m2 -startA - startB + 1 >= k:
                return self.findKthSortedArrays(nums1, startA, lenA, nums2, startB, lenB/2, k)  #delete the latter part of B
            else:
                return self.findKthSortedArrays(nums1, startA+lenA/2+1, lenA-lenA/2-1, nums2, startB, lenB, k-lenA/2-1)  #delete the start part of A
                

九章的写法,非常简洁但不够直接.

class Solution(object):
    def findMedianSortedArrays(self, nums1, nums2):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float
        """
        m = len(nums1) + len(nums2)
        if m%2 ==1 :
            return self.findKth(nums1, nums2, 0, 0, m/2 + 1) #总共奇数个数
        else:
            return (self.findKth(nums1, nums2, 0, 0, m/2) + self.findKth(nums1, nums2, 0, 0, m/2+1))/2.0 #总共偶数个数
            
            
    def findKth(self, nums1, nums2, A_start, B_start, k):

        if A_start >= len(nums1):
            return nums2[B_start+k-1] #如果A数组中已经没有数,则全部在B数组中找,直接返回这个元素
        if B_start >= len(nums2):
            return nums1[A_start+k-1] #如果B数组中已经没有数,则全部在A数组中找,直接返回这个元素
        if k == 1:
            return min(nums1[A_start], nums2[B_start])
        A_key = nums1[A_start + k/2-1] if (A_start+k/2-1) < len(nums1) else sys.maxint
        B_key = nums2[B_start + k/2 -1] if (B_start+k/2-1) < len(nums2) else sys.maxint
            
        if A_key < B_key:
            return self.findKth(nums1, nums2, A_start + k/2, B_start, k-k/2)
        else:
            return self.findKth(nums1, nums2, A_start, B_start + k/2, k-k/2)