LeetCode题解4：Median of Two Sorted

Median of Two Sorted Arrays问题：求2个有序数组的中位数，要求算法时间复杂度为O(log(m+n))。

难度：困难

思路：此题为“在2个有序数组中寻找第k个元素”问题的变体。根据中位数的定义，此题可转化为“在2个有序数组中寻找第‘(m+n)/2’个元素”。算法难点在于要求时间复杂度为O(log(m+n))，因此需要充分利用2个数组有序这一条件，每次循环将搜索空间缩小一半。

陷阱：数组下标计算。另外由于LeetCode网站测试用例设置问题，线性时间复杂度的算法也可能被接受。

代码：

class Solution:
    # @param {integer[]} nums1
    # @param {integer[]} nums2
    # @return {float}
    def find(self, nums1, nums2, k):
        if not nums1:
            return nums2[k]
        if not nums2:
            return nums1[k]
        ia, ib = len(nums1) // 2 , len(nums2) // 2
        ma, mb = nums1[ia], nums2[ib]
        if ia + ib < k:
            if ma > mb:
                return self.find(nums1, nums2[ib + 1:], k - ib - 1)
            else:
                return self.find(nums1[ia + 1:], nums2, k - ia - 1)
        else:
            if ma > mb:
                return self.find(nums1[:ia], nums2, k)
            else:
                return self.find(nums1, nums2[:ib], k)

    def findMedianSortedArrays(self, nums1, nums2):
        l = len(nums1) + len(nums2)
        if l % 2 == 1:
            return self.find(nums1, nums2, l // 2)
        else:
            return (self.find(nums1, nums2, l // 2) + self.find(nums1, nums2, l // 2 - 1)) / 2.0

另外，有人指出Python数组分片操作本身具有线性时间复杂度，因此严格地说上述算法不具备O(log(m+n))的时间复杂度。改进方法是将数组起止下标作为变量传给find函数，代码如下：

class Solution:
    # @param {integer[]} nums1
    # @param {integer[]} nums2
    # @return {float}
    def find(self, nums1, s1, e1, nums2, s2, e2, k):
        if e1 - s1 < 0:
            return nums2[k + s2]
        if e2 - s2 < 0:
            return nums1[k + s1]
        if k < 1:
            return min(nums1[k + s1], nums2[k + s2])
        ia, ib = (s1 + e1) // 2 , (s2 + e2) // 2
        ma, mb = nums1[ia], nums2[ib]
        if (ia - s1) + (ib - s2) < k:
            if ma > mb:
                return self.find(nums1, s1, e1, nums2, ib + 1, e2, k - (ib - s2) - 1)
            else:
                return self.find(nums1, ia + 1, e1, nums2, s2, e2, k - (ia - s1) - 1)
        else:
            if ma > mb:
                return self.find(nums1, s1, ia - 1, nums2, s2, e2, k)
            else:
                return self.find(nums1, s1, e1, nums2, s2, ib - 1, k)

    def findMedianSortedArrays(self, nums1, nums2):
        l = len(nums1) + len(nums2)
        if l % 2 == 1:
            return self.find(nums1, 0, len(nums1) - 1, nums2, 0, len(nums2) - 1, l // 2)
        else:
            return (self.find(nums1, 0, len(nums1) - 1, nums2, 0, len(nums2) - 1, l // 2) + self.find(nums1, 0, len(nums1) - 1, nums2, 0, len(nums2) - 1, l // 2 - 1)) / 2.0

附：线性时间复杂度算法代码：

class Solution:
    # @param {integer[]} nums1
    # @param {integer[]} nums2
    # @return {float}
    def findMedianSortedArrays(self, nums1, nums2):
        mid = (len(nums1) + len(nums2) >> 1) + 1
        i = j = 0
        median = n = None
        for k in xrange(mid):
            v1 = nums1[i] if i < len(nums1) else None
            v2 = nums2[j] if j < len(nums2) else None
            n = median
            if v2 is None or (v1 is not None and v1 < v2):
                median = v1
                i += 1
            else:
                median = v2
                j += 1
        if (len(nums1) + len(nums2)) % 2 == 0 and n is not None:
            return (median + n) / 2.0
        return median