[leetcode] 寻找两个有序数组的中位数

题目描述

给定两个大小为 m 和 n 的有序数组 nums1 和 nums2。
请你找出这两个有序数组的中位数，并且要求算法的时间复杂度为 O(log(m + n))。
你可以假设 nums1 和 nums2 不会同时为空。

输入示例:

示例 1:
nums1 = [1, 3]
nums2 = [2]
则中位数是 2.0

示例 2:
nums1 = [1, 2]
nums2 = [3, 4]
则中位数是 (2 + 3)/2 = 2.5

思路1

直接将两个数组执行归并排序合并为一个数组nums，然后便可以得到中位数。
时间复杂度： $o(n)$
空间复杂度： $o(n)$

代码

def findMedianSortedArrays(self, nums1, nums2):
    """
    先进行归并排序，然后直接取中位数
    :param nums1:
    :param nums2:
    :return:
    """
    nums = []
    i = j = 0
    while i < len(nums1):
        n1 = nums1[i]
        while j < len(nums2):
            n2 = nums2[j]
            if n2 < n1:
                nums.append(n2)
                j += 1
            else:
                break
        nums.append(n1)
        i += 1
    if j < len(nums2):
        nums.extend(nums2[j:])

    len_nums = len(nums)
    if len_nums % 2 == 1:
        return nums[(len_nums - 1) // 2]
    else:
        return float(nums[len_nums // 2 - 1] + nums[len_nums // 2]) / 2

思路2

这种方法比思路1要复杂很多，需要有两个定理支持：

定理一：从有序序列 $seq1$ 的首尾区域（不包含中值）内随机剥离等数目的数值得到新的有序序列 $seq2$ ， $seq1$ 与 $seq2$ 的中位数相同。
定理一示例

定理二：对于有序序列 $seq1$ 、 $seq2$ 及其中位数 $med1$ 、 $med2$ ，当从中位数更小的序列的左侧、中位数更大的序列的右侧删除等数目的元素后，新得到序列 $seq1'$ 和 $seq2'$ ， $seq1$ 合并 $seq2$ 的中位数与 $seq1'$ 合并 $seq2'$ 的中位数相同（基于定理一）。而当 $med1$ 与 $med2$ 相同时， $med1$ 即为合并后的中位数。
定理二示例
基于上面两个定理，算法如下所述：

从nums1和nums2中分别算得中位数m1,m2
用两个中位数进行比较，取中间的两段——即较大中位数所在数组的左段，较小中位数所在数组的右段。左右两段形成时，需要从原本数组中剥离等数量的数值。
从新的两端数组中寻找中位数，形成递归
当其中一个数组长度较小时，直接快速得到中位数

时间复杂度： $o(n)$ 但理论上上比思路一要小
空间复杂度： $o(0)$

代码

def find_median(self, nums):
    len_nums = len(nums)
    if len_nums % 2 == 0:
        # 偶数返回中间两值的均值，同时index为右边
        return (nums[len_nums >> 1] + nums[(len_nums >> 1) - 1]) / 2., len_nums >> 1
    else:
        return nums[((len_nums + 1) >> 1) - 1], ((len_nums + 1) >> 1) - 1


def findMedianSortedArrays(self, nums1, nums2):
    len_nums1 = len(nums1)
    len_nums2 = len(nums2)
    # 递归终止，找到了中位数
    if len_nums1 == 0:  # 其中一个数组为空
        median, _ = self.find_median(nums2)
        return median
    if len_nums2 == 0:
        median, _ = self.find_median(nums1)
        return median
    if len_nums1 <= 2 or len_nums2 <= 2:  # 其中一个数组长度不足
        nums = []
        nums.extend(nums1)
        nums.extend(nums2)
        nums.sort()
        median, _ = self.find_median(nums)
        return median

    median1, med_idx1 = self.find_median(nums1)
    median2, med_idx2 = self.find_median(nums2)
    if median1 == median2:  # 两数组的中位数相等，直接返回
        return median1
    else:  # 不等进入递归
        if median1 > median2:  # 取nums1的左侧，nums2的右侧
            if len_nums2 % 2 == 0:  # 偏移偶数数组nums2的指针，从而子数组能够包含中间两位
                med_idx2 -= 1
            strip_len = min(len(nums1) - med_idx1 - 1,
                            med_idx2)  # 这里需要等额剥离相同长度。依据原理：从有序序列的首尾区域（不包含中值）内随机剥离等量的数值，最终的中位数不变
            new_num1 = nums1[0:len(nums1) - strip_len]
            new_num2 = nums2[-(len(nums2) - strip_len):]
        else:
            if len_nums1 % 2 == 0:
                med_idx1 -= 1
            strip_len = min(len(nums2) - med_idx2 - 1, med_idx1)
            new_num2 = nums2[0:len(nums2) - strip_len]
            new_num1 = nums1[-(len(nums1) - strip_len):]
        return self.findMedianSortedArrays(new_num1, new_num2)

参考

[1] leetcode原题
[2] 详尽源码