LeetCode题解 4. 寻找两个正序数组的中位数

Author：baiyucraft

BLog： baiyucraft’s Home

---

一、题目描述

题目地址：4. 寻找两个正序数组的中位数

  给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的 中位数 。

输入：nums1 = [1,3], nums2 = [2]

输出：2.00000

解释：合并数组 = [1,2,3] ，中位数 2

输入：nums1 = [1,2], nums2 = [3,4]

输出：2.50000

解释：合并数组 = [1,2,3,4] ，中位数 (2 + 3) / 2 = 2.5

输入：nums1 = [0,0], nums2 = [0,0]

输出：0.00000

输入：nums1 = [], nums2 = [1]

输出：1.00000

输入：nums1 = [2], nums2 = []

输出：2.00000

二、思路以及代码

  此题一拿到手，感觉有个非常朴素的最暴力的解法，先合并两个数组，然后排序找出中位数，对于Python来说，代码实现无比简单，所以第一种解法就是暴力解法。

1.暴力解法

代码：

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        nums1 += nums2
        nums1.sort()
        n = len(nums1)
        mid = n >> 1
        if n & 1:
            return nums1[mid]
        else:
            return (nums1[mid] + nums1[mid - 1]) / 2

复杂度分析：

时间复杂度：$o(n{\log }_{2}n)$，sort()方法是归并排序

空间复杂度：$o(n)$

2.双指针分别遍历

思路： 在第一种方法中，不仅要合并数组，还要对数组进行排序，那么很自然的想到，是否可以维护两个指针分别遍历两个数组，然后根据计算出的中位数以及两个指针所对应的元素大小来找到中位数。

  我们可以依照这种想法写出大致的代码：

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m = len(nums1)
        n = len(nums2)
        i, j, count = 0, 0, (m + n + 1) >> 1
        # 一般情况下遍历，找到起始下标
        while count:
            if nums1[i] < nums2[j]:
                i += 1
                count -= 1
            else:
                j += 1
                count -= 1
         # 奇数
        if (m + n) & 1:
            return max(nums1[i - 1], nums2[j - 1])
        # 偶数
        else:
            return (max(nums1[i - 1], nums2[j - 1]) + min(nums1[i], nums2[j])) / 2.0

  代码中对于最后返回值的解释：

• 若总长度m + n为奇数，由于不知道最后一个遍历到的是nums1[i - 1]还是nums2[j - 1]，但由于升序的关系，最后遍历的值一定是已经遍历过的中最大的，所以取两者最大，即max(nums1[i - 1], nums2[j - 1])
• 若总长度m + n为偶数，因为要找到两个数取平均值，通过代码可以知道，经过(m + n + 1) >> 1次遍历后的最后一个数，是中间偏左的那个数mid_left，而下一个继续遍历的应该是中间偏右的那个数mid_right。按照前面奇数时的思想，由此可得两个数分别是：
  • mid_left是已经遍历的数中最大的，为max(nums1[i - 1], nums2[j - 1])
  • mid_right是未遍历的数中最小的，为min(nums1[i], nums2[j])

  但是这里可以发现，若是nums2中的数，比nums1中所有数还要大，例如：

nums1 = [1, 2, 3, 4]

nums2 = [5]

  这种情况下，count为3，结束循环时i = 3, j = 0。此时，若是执行max(nums1[i - 1], nums2[j - 1])，则j - 1 = -1会超出列表索引 。

  还有一种情况，若nums1或num2遍历完了，例如：

nums1 = [1, 2]

nums2 = [3, 4, 5, 6]

  这种情况下，结束循环时i = 2, j = 1。此时，若是执行min(nums1[i], nums2[j]) ，则i = 2会超出列表索引。

  针对以上的情况，这里给出一种解决办法，在遍历数组前进行预处理，在数组最前面前面插入一个最小值-2147483648，在数组最后面插入一个最大值2147483647，并且从i = 1, j = 1开始遍历，处理后的数组示意如下：

new_nums = [-2147483648, nums[0], nums[1], ······, nums[n - 1], 2147483647]

动画：

代码：

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m = len(nums1)
        n = len(nums2)
        # 0:-inf 1-m:[] m+1:inf
        nums1.insert(0, -2147483648)
        nums1.append(2147483647)
        nums2.insert(0, -2147483648)
        nums2.append(2147483647)
        # 初始化
        i, j, count = 1, 1, (m + n + 1) >> 1
        # 一般情况下遍历，找到起始下标
        while count:
            if nums1[i] < nums2[j]:
                i += 1
                count -= 1
            else:
                j += 1
                count -= 1
        # 奇数
        if (m + n) & 1:
            return max(nums1[i - 1], nums2[j - 1])
        # 偶数
        else:
            return (max(nums1[i - 1], nums2[j - 1]) + min(nums1[i], nums2[j])) / 2.0

复杂度分析：

时间复杂度：$o(\frac{m+n}{2})=o(n)$

空间复杂度：$o(1)$，预处理了数组

3.二分查找

思路： 在第二种的方法中，我们已经通过双指针分别遍历的方式来找到了中位数，但是显然，因为是顺序查找的，所以查找的时间复杂度并不算高，题目的进阶是让我们设计一个时间复杂度为$o(lo{g}_{2}n)$的算法，那在查找的算法中，对于有序的数组，不免会想到二分查找。

  那么一个升序数组的二分查找大家都很熟悉，无非就是找到中点，如果target的值大于nums[mid]，则留下nums[mid]右边的数组继续查找，反之则nums[mid]查找左边的数组。

  对于本题，首先我们要明确一个目标，要通过二分查找到第k小的小的数，对于两个数组nums1和nums2，要找到第k小的数，我们先找到两个数组中第k/2 - 1的数，然后比较这两个数的大小。对于nums1[k/2 − 1]和nums2[k/2 − 1] 中的较小值，因为这两个数前面分别有下标0···(k/2 - 2)个数，即最多只会有(k/2 - 1) + (k/2 - 1) =k − 2个元素比它小，那么它就不能是第 k小的数了。再根据所有情况有以下讨论：

• 若nums[k/2 - 1] > nums2[k/2 - 1]，说明nums2[k/2 - 1]不是第k小的数，那nums2[0]到nums2[k/2 - 1]都不可能是第k小的数，排除。
• 若nums[k/2 - 1] < nums2[k/2 - 1]，同上排除nums1[0]到nums1[k/2 - 1]
• 若nums[k/2 - 1] = nums2[k/2 - 1]，同第一点

  排除了数组后，也就相当于缩小了范围，在新建立的数组上继续用此法查找k - (k/2 -1) = k/2 + 1小的数，最后找到第k小的数。其中需要处理特殊情况：

• 如果nums1或nums2的长度小于k/2 - 1，会造成索引越界，所以比较的是该数组的最后一个元素，下一次遍历的也应当是第k - len(nums)小的元素。
• 如果一个数组为空，我们可以直接返回另一个数组中第k小的元素。
• 如果k = 1，我们只要返回两个数组首元素的最小值即可。

动画：

代码：

# 递归
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m = len(nums1)
        n = len(nums2)
        k= (m + n + 1) >> 1

        # 找第k小的数
        def get_kth(nums1, nums2, k):
            # 保持nums1比较长
            if len(nums1) < len(nums2):
                nums1, nums2 = nums2, nums1
            # 短数组空，直接返回
            if len(nums2) == 0:
                return nums1[k - 1]  
            # 若为1，直接比较数组首位
            if k == 1:
                return min(nums1[0], nums2[0])
            
            tmp = min(k >> 1, len(nums2))  # 保证不上溢
            if nums1[tmp - 1] >= nums2[tmp - 1]:
                return get_kth(nums1, nums2[tmp:], k - tmp)
            else:
                return get_kth(nums1[tmp:], nums2, k - tmp)

        if (m + n) & 1:
            return get_kth(nums1, nums2, k)
        else:
            return (get_kth(nums1, nums2, k) + get_kth(nums1, nums2, k + 1)) / 2.0

# 循环，参考官方题解
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m = len(nums1)
        n = len(nums2)
        k = (m + n + 1) >> 1

        def get_kth(k):
            p1, p2 = 0, 0
            while True:
                # nums1 为空
                if p1 == m:
                    return nums2[p2 + k - 1]
                # nums2 为空
                if p2 == n:
                    return nums1[p1 + k - 1]
                # 若为1，直接比较数组首位
                if k == 1:
                    return min(nums1[p1], nums2[p2])

                tmp = (k >> 1) - 1
                newp1 = min(p1 + tmp, m - 1)
                newp2 = min(p2 + tmp, n - 1)
                if nums1[newp1] <= nums2[newp2]:
                    k = k - (newp1 - p1 + 1)
                    p1 = newp1 + 1
                else:
                    k = k - (newp2 - p2 + 1)
                    p2 = newp2 + 1

        if (m + n) & 1:
            return get_kth(k)
        else:
            return (get_kth(k) + get_kth(k + 1)) / 2.0

复杂度分析：

时间复杂度：$o(lo{g}_2(m+n))$

空间复杂度：$o(1)$

4.划分数组

**思路：**这种方法参考了官方题解，叙述一遍推导过程

  首先，我们得知道中位数的作用：

中位数将一个集合划分成两个子集，其中一个子集中的元素总是大于另一个子集中的元素

  这里先对对数组作第二种方法一样的进行预处理，即在数组最前面前面插入一个最小值-2147483648，在数组最后面插入一个最大值2147483647，具体原因后文再讲，即：

new_nums = [-2147483648, nums[0], nums[1], ······, nums[n - 1], 2147483647]

new_nums[0] = -inf，new_nums[m] = inf

  知道了中位数的作用，也就是说要找到两个数组的分割点i和j，使分割后的nums1_right和nums2_right组成的集合left_right中的元素总是大于nums1_left和nums2_left组成的集合right_part中的元素，也就是说划分后的i和j应该满足以下条件：

• 总长度m + n为奇数时，len(left_part) = len(right_part) + 1
• 总长度m + n为偶数时，len(left_part) = len(right_part)
• max(left_part) <= min(right_part)

  要满足以上条件，针对如下划分：

			      left_part	       	   	|				right_part
 nums1[0], nums1[1], ..., nums1[i - 1]	|  nums1[i], nums1[i + 1], ..., nums1[m - 1]
 nums2[0], nums2[1], ..., nums2[j - 1]	|  nums2[j], nums2[j + 1], ..., nums2[n - 1]

  只需要保证：

• 针对前两条，i + j = m - i + n - j + 1（奇数）或i + j = m - i + n - j（偶数），即i + j = (m + n + 1) // 2，为了保证统一，让len(nums1) <= len(nums2)，即m <= n。这样就可以保证针对每一个i都有唯一且存在的j = (m + n + 1) // 2 - i对应

• 针对第三条，只需要满足划分后，nums1[i - 1] <= nums2[j]以及nums2[j - 1] <= nums1[i]

  • 为了使索引不越界，对数组做了处理

  • 针对预处理后的数组，我们所需要做的是，在索引为1 ~ m-1中找到i，使得：nums1[i - 1] <= nums2[j]且nums2[j - 1] <= nums1[i]

  • 上述式子等价为，在索引为1 ~ m-1中找到最大的i，使得：nums1[i - 1] <= nums2[j]，

  • 当i从1 ~ m-1递增时，nums1[i − 1]递增，nums2[j]递减，所以一定存在一个最大的i满足nums1[i − 1] <= nums2[j]

    如果i是最大的，那么说明i + 1不满足。将i + 1带入可以得到nums1[i] > nums2[j - 1]，也就是nums2[j - 1] <= nums1[i]

  因此我们可以对i在1 ~ m-1中进行二分搜索，找打最大的满足nums1[i − 1] <= nums2[j]的值，此时的i和j就是正确的划分方法。

动画：

代码：

    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        # 预处理 
        # 保证 nums1 为较短的
        if len(nums1) > len(nums2):
            nums1, nums2 = nums2, nums1
        # 0:-inf 1-m:[] m+1:inf
        nums1 = [-2147483648] + nums1 + [2147483647]
        nums2 = [-2147483648] + nums2 + [2147483647]

        m = len(nums1)
        n = len(nums2)
        # 二分查找起始位置
        left, right = 1, m - 2
        # 在循环中常用的一部分
        len2 = (m + n + 1) >> 1
        # median1：前一部分的最大值, median2：后一部分的最小值
        median1, median2 = 0, 0

        while left <= right:
            # 前一部分包含 nums1[0 .. i-1] 和 nums2[0 .. j-1]
            # 后一部分包含 nums1[i .. m-1] 和 nums2[j .. n-1]
            i = (left + right) >> 1
            j = len2 - i
			
            if nums1[i - 1] <= nums2[j]:
                median1, median2 = max(nums1[i - 1], nums2[j - 1]), min(nums1[i], nums2[j])
                left = i + 1
            else:
                right = i - 1

        if (m + n) & 1:
            return median1
        else:
            return (median1 + median2) / 2.0

复杂度分析：

时间复杂度：$o(lo{g}_{2}min(m,n))$

空间复杂度：$o(1)$