LeetCode-4 寻找两个有序数组的中位数

发表于 | 更新于 | 分类于 | 评论数： | 阅读次数：

题目：寻找两个有序数组的中位数

给定两个大小为 m 和 n 的有序数组 nums1 和 nums2。

请你找出这两个有序数组的中位数，并且要求算法的时间复杂度为 O(log(m + n))。

你可以假设 nums1 和 nums2 不会同时为空。

示例 1:

1
2
3
4
nums1 = [1, 3]
nums2 = [2]

则中位数是 2.0

示例 2:

1
2
3
4
nums1 = [1, 2]
nums2 = [3, 4]

则中位数是 (2 + 3)/2 = 2.5

思路

这题目乍一看有点绕，其实简单一点考虑就是找这两个数组集合的中位数。

  • 先把两个数组合并成一个数组temp = nums1 + nums2
  • 然后再对这个新数组进行排序
  • 找中位数：
    • 当这个数组长度是奇数时：中位数就是排在最中间的那个数
    • 当这个数组长度是偶数时：中位数就是排在最中间的两个数的平均数

实现如下：

1
2
3
4
5
6
7
8
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        temp = nums1 + nums2
        temp.sort()
        if len(temp) % 2 ==1:
            return temp[int(len(temp) / 2)]
        else:
            return (temp[int(len(temp) / 2) - 1] + temp[int(len(temp) / 2)]) / 2.0

官方解答

这题的官方解答看起来比较复杂，但考虑到上面的解法用到了sort()方法，可能本身就存在算法复杂度高的问题。所以官方解答使用了递归排序的方法。

  • 先不考虑特殊情况，分别将两个数组划分成两部分

    1
    2
    3
          left_part          |        right_part
    A[0], A[1], ..., A[i-1]  |  A[i], A[i+1], ..., A[m-1]
    B[0], B[1], ..., B[j-1]  |  B[j], B[j+1], ..., B[n-1]

  • 那么如果left_part的长度等于right_part的长度，且right_part的值都大于left_part的值（即min(right_part)>=max(left_part)）。我们就找到了中位数：

    $median=\frac{max(left_part)+min(right_part)}{2}​$

  • 那么i和j有什么关系呢？

    • 当总长度$m+n​$为偶数时：

      两边长度相等：$i+j=m-i+n-j$ ，所以可以得到$j=\frac{m+n}{2}-i$。

    • 当总长度$m+n​$为奇数时：

      左边比右边多1个数：$i+j=m-i+n-j+1$，所以得到：$j=\frac{m+n+1}{2}-i$。

      由于$i, j$会不断移动，我们取较大值：$j=\frac{m+n+1}{2}-i$。

  • 先不考虑临界情况，并假设$m<n$，我们需要在$[0,m]$中找到$i$可以使$B[j-1]<=A[i]$且$A[i-1]<=B[j]$。

  • 现在我们可以开始按二叉树进行查找：

    • 初始化i的查找范围$imin = 0​$,$ imax = m​$,$i=\frac{imin+imax}{2}​$。并动态调整$imin,imax​$的值,为了方便可以不需要考虑奇数偶数，通过int()取整即可，如果是奇数那么i与原值一样，会继续调整$imin,imax​$的值，直到下一个
    • 当$i<m, B[j-1]>A[i]$时，我们需要增加$imin = imin + 1$
    • 当$i>0, A[i-1]>B[j]$时，我们需要减小$imax = imax -1$。

  • 最终结果，要考虑下特殊情况：

    • 左边最大值max_of_left：

      • 当$i=0$时：$B[j-1]$
      • 当$j=0$时：$A[i-1]$

      • max(a[i-1], B[j-1])

        当总长度m+n为奇数时，直接返回左边最大值

    • 右边最小值min_of_right:

      • 当$i=m$时：$B[j]$
      • 当$j=0$时：$A[i]$

      • Min(A[i], B[j])

        当总长度为偶数时，返回（max_of_left + min_of_right）/ 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
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Solution:
    def median(self, nums1, nums2):
        # Solution 2
        m, n = len(nums1), len(nums2)
        if m > n:
            nums1, nums2, m, n = nums2, nums1, n, m
        if n == 0:
            raise ValueError
        imin, imax, half_len = 0, m, int((m + n + 1) / 2)
        while imin <= imax:
            i = int((imin + imax) / 2)
            j = half_len - i
            if i < m and nums2[j - 1] > nums1[i]:
                # i is too small, must increase it
                imin = i + 1
            elif i > 0 and nums1[i - 1] > nums2[j]:
                # i is too big, must decrease it
                imax = i - 1
            else:
                # i is perfect
                if i == 0:
                    max_of_left = nums2[j - 1]
                elif j == 0:
                    max_of_left = nums1[i - 1]
                else:
                    max_of_left = max(nums1[i - 1], nums2[j - 1])

                if (m + n) % 2 == 1:
                    return max_of_left

                if i == m:
                    min_of_right = nums2[j]
                elif j == n:
                    min_of_right = nums1[i]
                else:
                    min_of_right = min(nums1[i], nums2[j])

                return (max_of_left + min_of_right) / 2.0

还是有点绕，如果大家有更好的方法或者思路，欢迎一起探讨。