Median of Two Sorted Arrays

2016-05-01 本文已影响241人 M23

Median of Two Sorted Arrays

这是一个leetcode上的算法题目，标记为hard。具体描述如下：

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

leetcode上要求实现接口如下：

class Solution(object):
def findMedianSortedArrays(self, nums1, nums2):
"""
:type nums1: List[int]
:type nums2: List[int]
:rtype: float
"""

在做题之前，首先要明白什么是中位数。以下是来自某度的解释：

中位数（又称中值，英语：Median），统计学中的专有名词，代表一个样本、种群或概率分布中的一个数值，其可将数值集合划分为相等的上下两部分。对于有限的数集，可以通过把所有观察值高低排序后找出正中间的一个作为中位数。如果观察值有偶数个，通常取最中间的两个数值的平均数作为中位数。

举个栗子，数列[1,2,2,3,3,4]的中位数为2.5。两个序列的中位数则是将两个数组merge到同一个序列中，然后取中位数。栗子又来了[1,3,7,8]和[2,4,5,6]的中位数是4.5。
看到算法中提到的时间复杂度为log(m+n)，很明显，这里需要二分搜索。

二分搜索的挑战

Knuth在其鸿篇巨制The Art of Computer Programming, volume3，Sorting and Searching的6.2.1节曾指出，虽然第一篇二分搜索论文在1946年就发表了，但第一个没有错误的二分搜索程序却直到1962年才出现。所以，的确很难。继续阅读下文之前，各位读者不妨先拿出纸笔，花几分钟粗略设计一下这个题目的算法。

times-up.jpg

设计实现

回到题目，设两个有序数组为A和B，长度分别为m、n，如何才能最快的找到中位数呢。不失一般性，可以假定m 小于等于 n,原因是A和B的中位数必然等于B和A的中位数，原因是AB的顺序并不影响AB组合后的序列，因此得证。

确定边界

二分查找，每一步都需要缩小搜索的范围，那么不难想到，本题一个可以缩小搜索范围的做法是通过比较两个数组各自的中位数m1,m2，有以下三种情况来区分（数组中冒号是借用python切片表达）：

m1 == m2, 中大奖，直接返回m1(或m2）
m1 < m2, 返回数组A[?:?]和B[?:?] 的中位数
m1 > m2, 返回数组A[?:?] 和 B[?:?]的中位数

上述列表中的边界中有诸多问号，这也是二分查找的关键之一。我们通过分析分别填上，考虑下面两个因素：

中位数在序列包含元素的奇偶性上表现不同：如果序列元素个数为奇数作为中位数是数组中的某数，否则是两个数的平均值。所以在处理边界上，也是和奇偶相关的。
范围缩小的一致性：这里指的并不是等比例，而是具体的数字。即当数组A减少了n个元素时，数组B也必须减少n个，否则结果肯定是不对的，试想A=[2,2], B=[1,2,3...8,9]。
基于这两点，

        if nums1 and nums2:
            m1 = findMedianOfSingleSortedArray(nums1)
            m2 = findMedianOfSingleSortedArray(nums2)
            if m1 == m2:
                return m1
            if m1 < m2:
                if len1 % 2 == 0:
                    return self.findMedianSortedArrays(nums1[len1 / 2 - 1:], nums2[:len2 - len1 / 2 + 1])
                return self.findMedianSortedArrays(nums1[len1 / 2:], nums2[: len2 - len1 / 2])
            else:
                if len1 % 2 == 0:
                    return self.findMedianSortedArrays(nums1[:len1 / 2 + 1 ], nums2[len1 / 2 - 1:])
                return self.findMedianSortedArrays(nums1[:len1 / 2 + 1 ], nums2[len1 / 2:])

边界确定后，下一步就是要找到结束条件了。

结束条件

开始之前，还是惯例，希望读者能先思考一下。何时终止二分查找。

同时，这里先插播一个有意思的感觉：做英语选择题拿不太准时，比如第一感觉选A，然后修改了C，结果改错了。然后英语老师强调第一感觉很重要 (其实颇有点孕妇效应的感觉)。数学的题目第一感觉选B，后来仔细推理，原来D才是正确答案，数学老师强调的是千万别信第一感觉。当然不排除邪恶的数学出题人故意挖坑给大家。回到这个题目，其实的确也是个数学题，小坑呢，也是有的：
当范围逐渐缩小，是否最终是其中一个数组变为空，然后计算另外一个数组的中位数就可以了呢？估计我不说你也能猜出来，这个想法是错的。原因就在于，有可能中位数包含在了逐渐缩小的范围中，尤其是在最后变为空的之前一段时间。
如果不为空，那各个数组包含多少元素时应该结束呢。仔细一想，这取决与最后的中位数到底需要几个数字才能算出来，答案是，当序列总数是偶数时，需要2个，奇数时，需要1个。尝试解释一下原因：
假设[2,3]和[0,1,5,6]这种情况，[2,3]是不可再缩减的，因为一旦再缩小范围，那么中位数就无法得出了。同理，可得出奇数时需要1个的结论。
那么结束条件，也不难得出,

        if len1 == 1:
            if len2 % 2 == 0:
                return least_1_even(nums1, nums2)
            else:
                return least_1_odd(nums1, nums2)

        if len1 == 2:
            if len2 % 2 == 0:
                return least_2_even(nums1, nums2)
            else:
                return least_2_odd(nums1, nums2)

简单解释一下，为何只需要判断第一个数组长度是否到达了1和2。因为之前假定A的长度小于B，并且我们在缩小范围时，总是缩小相同的数目。所以必然是A先达到1、2。不需要关注B的长度。一个数组长度为1或2，另外一个数组长度为奇数或偶数需要不同处理，这里都是苦力活，即比较A中的1个或2个数字与B的中位数的大小关系。具体的参考代码，逻辑比较简单，但需要非常仔细。千万不要遗漏任何一种情况。

def least_2_even(nums1, nums2):
    a, b = nums1[0],nums1[1]
    len2 = len(nums2)
    p, q = nums2[len2 / 2 - 1], nums2[len2 / 2]
    if b <= q and p <= a:
        return (a + b) / 2.
    if b <= p:
        if len2 > 2:
            return (max(b, nums2[len2 / 2 - 2]) + p) / 2.
        return (b + p) / 2.
    if a >= q:
        if len2 > 2:
            return (min(a, nums2[len2 / 2 + 1]) + q) / 2.
        return (a + q) / 2.
    if a <=p and q <=b:
        return (p + q) / 2.
    if p <= a <= q <= b:
        return (a + q) / 2.
    return (b + p) / 2.
def least_2_odd(nums1, nums2):
    a, b = nums1[0],nums1[1]
    len2 = len(nums2)
    m = nums2[len2 / 2]
    if a <= m <= b:
        return m
    if b <= m:
        if len2 > 1:
            return max(b, nums2[len2 / 2 - 1])
        return b
    if m <= a:
        if len2 > 1:
            return min(a, nums2[len2 / 2 + 1])
        return a

def least_1_odd(nums1, nums2):
    a = nums1[0]
    len2 = len(nums2)
    m = nums2[len2 / 2]
    if len2 > 1:
        if a >= nums2[len2 / 2]:
            return (min(a, nums2[len2 / 2 + 1]) + m) / 2.
        else:
            return (max(a, nums2[len2 / 2 - 1]) + m) / 2.
    else:
        return (a + m) / 2.

def least_1_even(nums1, nums2):
    a = nums1[0]
    len2 = len(nums2)
    p, q = nums2[len2 / 2 - 1], nums2[len2 / 2]
    if p <= a <= q:
        return a
    if a < p:
        return p
    return q

例外

先别急着提交代码，还有一种异常情况需要处理呢。

A or B 为空：直接返回不空的数组的中位数
A and B都为空，题目不存在这种情况，忽略
代码：

        if not nums1:
            return findMedianOfSingleSortedArray(nums2)
        if not nums2:
            return findMedianOfSingleSortedArray(nums1)

附上工具函数：

def findMedianOfSingleSortedArray(l):
    ''' Array l must not be []'''
    length = len(l)
    if length % 2 == 0:
        return (l[length/2] + l[length /2 - 1])/2.
    return l[length/2]

OK，完成。

时间复杂度分析

对于数组A来讲，算法在每一步都能够缩短搜索范围一半，直至元素个数为2。这需要log(m)次操作，当两个数组元素分别个数为2和n-m+2时，需要常数次比较即可获得中位数，因此该算法时间复杂度为O(log(min(m,n))),是快于题目给出的O(log(m+n))的。（对于python的切片操作，可以有等价的传递数组index的O(1)的方式替代，故此处省略不记。）

总结

正如前面提到的，二分搜索的代码其实坑很多，整数除2到底是取上限还是下限，类似off-by-one则是每次写都会碰到。但我们只要时刻铭记二分搜索中的三个部分，就能够强迫自己保持清醒了：

初始化: 循环不变式(loop invariant)始终为真
保持: 如果在某次迭代开始时以及循环执行时，不变式都为真，那么循环执行完毕不变式依然为真
终止：循环能够终止，并且得到期望结果。
所有代码在我的github上

Median of Two Sorted Arrays