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)