吴良超的学习笔记

LeetCode解题报告(4)--二分法查找两个有序数组的中位数

发表于 | ，

原题如下：

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)).

大意就是找出两个排好序的数组中所有数的中位数。

最直观的做法是先将两个数组按从大到小合并成一个数组，再找出中位数。这样相当于遍历了两个数组，其时间复杂度是O(m+n)。实现代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Solution(object):
    def findMedianSortedArrays(self, nums1, nums2):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float
        """
        index1=0
        index2=0
        mergedArr=[]
        while index1<len(nums1) and index2<len(nums2):
            if nums1[index1] > nums2[index2]:
                mergedArr.append(nums2[index2])
                index2+=1
            elif nums1[index1] < nums2[index2]:
                mergedArr.append(nums1[index1])
                index1+=1
            else: #equal
                mergedArr.append(nums1[index1])
                mergedArr.append(nums2[index2])
                index1+=1
                index2+=1
            
        if index1!=len(nums1):
            while index1<len(nums1):
                mergedArr.append(nums1[index1])
                index1+=1
        
        if index2!=len(nums2):
            while index2<len(nums2):
                mergedArr.append(nums2[index2])
                index2+=1
        
        mergedLen=len(mergedArr)
        if mergedLen %2 == 0: #even length
            return float(mergedArr[mergedLen/2]+mergedArr[mergedLen/2-1])/2 
        else:
            return float (mergedArr[(mergedLen-1)/2])

尽管上面的时间复杂度是O(m+n)，但是实际的提交时也能AC。

虽然上面的方法也能够AC，但是根据题目的提示的O(log(m+n))时间复杂度，显然还有更好的解法。下面就来讨论一种时间复杂度为O(log(m+n))的算法。

提到时间复杂度为O(log(m+n))的算法，很容易想到的就是二分查找，所以现在要做的就是在两个排序数组中进行二分查找。

具体思路如下，可以将问题转化为在两个数组中找第K个大的数，先在两个数组中分别找出第k/2大的数，再比较这两个第k/2大的数，这里假设两个数组为A,B。那么比较结果会有下面几种情况：

  • A[k/2]=B[k/2],那么第k大的数就是A[k/2]
  • A[k/2]>B[k/2],那么第k大的数肯定在A[0:k/2+1]和B[k/2:]中，这样就将原来的所有数的总和减少到一半了，再在这个范围里面找第k/2大的数即可，这样也达到了二分查找的区别了。
  • A[k/2] < B[k/2]，那么第k大的数肯定在B[0:k/2+1]和A[k/2:]中,同理在这个范围找第k/2大的数就可以了。

上面思路的实现代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution:
    def findMedianSortedArrays(self, A, B):
        n = len(A) + len(B)
        if n % 2 == 1:
            return self.findKth(A, B, n / 2 + 1)
        else:
            smaller = self.findKth(A, B, n / 2)
            bigger = self.findKth(A, B, n / 2 + 1)
            return (smaller + bigger) / 2.0

    def findKth(self, A, B, k):
        if len(A) == 0:
            return B[k - 1]
        if len(B) == 0:
            return A[k - 1]
        if k == 1:
            return min(A[0], B[0])
        
        a = A[k / 2 - 1] if len(A) >= k / 2 else None
        b = B[k / 2 - 1] if len(B) >= k / 2 else None
        
        if b is None or (a is not None and a < b):
            return self.findKth(A[k / 2:], B[0:k/2+1], k - k / 2)
        return self.findKth(A[0:k/2+1], B[k / 2:], k - k / 2)