[Leetcode刷题] 第4题 Median of Two Sorted Arrays 寻找两个有序数组的中位数

题目：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)).

You may assume nums1 and nums2 cannot be both empty.

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

解题思路：二分法

    * 1.首先，让我们在任一位置 i 将 A(长度为m) 划分成两个部分：
    *            leftA            |                rightA
    *   A[0],A[1],...      A[i-1] |  A[i],A[i+1],...A[m - 1]
    *
    * 由于A有m个元素，所以有m + 1中划分方式(i = 0 ~ m)
    *
    * 我们知道len(leftA) = i, len(rightA) = m - i;
    * 注意：当i = 0时，leftA是空集，而当i = m时，rightA为空集。
    *
    * 2.采用同样的方式，将B也划分为两部分：
    *            leftB            |                rightB
    *   B[0],B[1],...      B[j-1] |   B[j],B[j+1],...B[n - 1]
    *  我们知道len(leftA) = j, len(rightA) = n - j;
    *
    *  将leftA和leftB放入一个集合，将rightA和rightB放入一个集合。再把这两个集合分别命名为leftPart和rightPart。
    *
    *            leftPart         |                rightPart
    *   A[0],A[1],...      A[i-1] |  A[i],A[i+1],...A[m - 1]
    *   B[0],B[1],...      B[j-1] |  B[j],B[j+1],...B[n - 1]
    *
    *   如果我们可以确认：
    *   1.len(leftPart) = len(rightPart); =====> 该条件在m+n为奇数时，该推理不成立
    *   2.max(leftPart) <= min(rightPart);
    *
    *   median = (max(leftPart) + min(rightPart)) / 2;  目标结果
    *
    *   要确保这两个条件满足：
    *   1.i + j = m - i + n - j(或m - i + n - j + 1)  如果n >= m。只需要使i = 0 ~ m，j = (m+n+1)/2-i =====> 该条件在m+n为奇数/偶数时，该推理都成立
    *   2.B[j] >= A[i-1] 并且 A[i] >= B[j-1]
    *
    *   注意:
    *   1.临界条件：i=0,j=0,i=m,j=n。需要考虑
    *   2.为什么n >= m ? 由于0 <= i <= m且j = (m+n+1)/2-i,必须确保j不能为负数。
    *
    *   按照以下步骤进行二叉树搜索
    *   1.设imin = 0,imax = m，然后开始在[imin,imax]中进行搜索
    *   2.令i = (imin+imax) / 2, j = (m+n+1)/2-i
    *   3.现在我们有len(leftPart) = len(rightPart)。而我们只会遇到三种情况：
    *
    *      ①.B[j] >= A[i-1] 并且 A[i] >= B[j-1]  满足条件
    *      ②.B[j-1] > A[i]。此时应该把i增大。 即imin = i + 1;
    *      ③.A[i-1] > B[j]。此时应该把i减小。 即imax = i - 1;

具体代码：

class Solution:
    def findMedianSortedArrays(self, A, B):
        l = len(A) + len(B)
        if l % 2 == 1: #判断两个列表总长度是奇数还是偶数
            return self.kth(A, B, l // 2) #是奇数
        else:
            return (self.kth(A, B, l // 2) + self.kth(A, B, l // 2 - 1)) / 2. #是偶数


    def kth(self, a, b, k): #k为合并后这个列表中median的位置
        if not a: # A为空，输出B的median
            return b[k]
        if not b: # B为空，输出A的median
            return a[k]
        ia, ib = len(a) // 2, len(b) // 2     # ia,ib分别为A的media的位置，B的median的位置
        ma, mb = a[ia], b[ib]  # ma,mb分别为A的media，B的median


        if ia + ib < k: #当k的值大于A,B的median的位置的值，A,B的median的位置的值的和还不够大，需要向后半部分搜索
            if ma > mb: #如果A的media大于B的median,B的前半部分就不包含k
                return self.kth(a, b[ib + 1:], k - ib - 1)
            else:    #如果A的media小于B的median，A的前半部分就不会包含k
                return self.kth(a[ia + 1:], b, k - ia - 1)

        else: #当k的值小于等于A,B的median的位置的值，A,B的median的位置的值的和够大了，需要向前半部分搜索
            if ma > mb:  #如果A的median大于B的，A的后半部分就不包含k
                return self.kth(a[:ia], b, k)
            else:     #如果A的median小于B的，B的后半部分就不包含k
                return self.kth(a, b[:ib], k)

if __name__ == '__main__':
    a = [1,2,2,2,3]
    b = [4,5,6,7,1,4,7]
    o = Solution()
    m=o.findMedianSortedArrays(A=a,B=b)
    print(m)