Stefan Pochmann 的上帝之手（6）寻找两个正序数组

2020-07-25 本文已影响0人 WilliamY

寻找两个正序数组的中位数

给定两个大小为 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

一、二分查找（22行）

从时间复杂度的要求 O(log(m + n))可知，该算法每次需要去掉一半候选数字，二分查找就符合要求。
链接：https://leetcode-cn.com/problems/median-of-two-sorted-arrays/solution/xun-zhao-liang-ge-you-xu-shu-zu-de-zhong-wei-s-114/
根据中位数的定义，当 $m+n$ 是奇数时，中位数是两个有序数组中的第 $(m+n)/2$ 个元素，当 $m+n$ 是偶数时，中位数是两个有序数组中的第 $(m+n)/2$ 个元素和第 $(m+n)/2+1$ 个元素的平均值。因此，这道题可以转化成寻找两个有序数组中的第 $k$ 小的数，其中 $k$ 为 $(m+n)/2$ 或 $(m+n)/2+1$ 。

假设两个有序数组分别是 $\text{A}$ 和 $\text{B}$ 。要找到第 $k$ 个元素，我们可以比较 $\text{A}[k/2-1]$ 和 $\text{B}[k/2-1]$ ，其中 $/$ 表示整数除法。由于 $\text{A}[k/2-1]$ 和 $\text{B}[k/2-1]$ 的前面分别有 $\text{A}[0\,..\,k/2-2]$ 和 $\text{B}[0\,..\,k/2-2]$ ，即 $k/2-1$ 个元素，对于 $\text{A}[k/2-1]$ 和 $\text{B}[k/2-1]$ 中的较小值，最多只会有 $(k/2-1)+(k/2-1) \leq k/2-2$ 个元素比它小，那么它就不能是第 $k$ 小的数了。

因此我们可以归纳出三种情况：

如果 $\text{A}[k/2-1] < \text{B}[k/2-1]$ ，则比 $\text{A}[k/2-1]$ 小的数最多只有 $\text{A}$ 的前 $k/2-1$ 个数和 $\text{B}$ 的前 $k/2-1$ 个数，即最多只有 $k-2$ 个，因此 $\text{A}[k/2-1]$ 不可能是第 $k$ 个数， $\text{A}[0]$ 到 $\text{A}[k/2-1]$ 也都不可能是第 $k$ 个数，可以全部排除。
如果 $\text{A}[k/2-1] > \text{B}[k/2-1]$ ，则可以排除 $\text{B}[0]$ 到 $\text{B}[k/2-1]$ 。
如果 $\text{A}[k/2-1] = \text{B}[k/2-1]$ ，则可以归入第一种情况处理。

力扣官网

可以看到，比较 $\text{A}[k/2-1]$ 和 $\text{B}[k/2-1]$ 之后，可以排除 $k/2$ 个不可能是第 $k$ 小的数，查找范围缩小了一半。同时，我们将在排除后的新数组上继续进行二分查找，并且根据我们排除数的个数，减少 $k$ 的值，这是因为我们排除的数都不大于第 $k$ 小的数。

有以下三种情况需要特殊处理：

如果 $\text{A}[k/2-1]$ 或者 $\text{B}[k/2-1]$ 越界，那么我们可以选取对应数组中的最后一个元素。在这种情况下，我们必须根据排除数的个数减少 $k$ 的值，而不能直接将 $k$ 减去 $k/2$ 。
如果一个数组为空，说明该数组中的所有元素都被排除，我们可以直接返回另一个数组中第 $k$ 小的元素。
如果 $k=1$ ，我们只要返回两个数组首元素的最小值即可。

def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
    m, n = len(nums1), len(nums2)
    def findKthElement(k):
        index1, index2 = 0, 0
        while True:
            if index1 > m - 1:
                return nums2[index2 + k - 1]
            if index2 > n - 1:
                return nums1[index1 + k - 1]
            if k == 1:
                return min(nums1[index1], nums2[index2])
            
            # 正常情况
            n1 = min(m - 1, index1 + k // 2 - 1)
            n2 = min(n - 1, index2 + k // 2 - 1)
            if nums1[n1] <= nums2[n2]:
                k -= n1 - index1 + 1
                index1 = n1 + 1
            else:
                k -= n2 - index2 + 1
                index2 = n2 + 1

    if (m + n) % 2:
        return findKthElement((m + n) // 2 + 1)
    else:
        return (findKthElement((m + n) // 2) + findKthElement((m + n) // 2 + 1)) / 2

划分数组（19行）

思路与算法
为了使用划分的方法解决这个问题，需要理解「中位数的作用是什么」。在统计中，中位数被用来：

将一个集合划分为两个长度相等的子集，其中一个子集中的元素总是大于另一个子集中的元素。

如果理解了中位数的划分作用，就很接近答案了。
首先，在任意位置 $i$ 将 $\text{A}$ 划分成两个部分：

           left_A            |          right_A
    A[0], A[1], ..., A[i-1]  |  A[i], A[i+1], ..., A[m-1]

由于 $\text{A}$ 中有 $m$ 个元素，所以有 $m+1$ 种划分的方法（ $i \in [0, m]$ ）。

$\text{len}(\text{left_A}) = i, \text{len}(\text{right_A}) = m - i$ .

注意：当 $i = 0$ 时， $\text{left_A}$ 为空集，而当 $i = m$ 时, $\text{right_A}$ 为空集。

采用同样的方式，在任意位置 $j$ 将 $\text{B}$ 划分成两个部分：

           left_B            |          right_B
    B[0], B[1], ..., B[j-1]  |  B[j], B[j+1], ..., B[n-1]

将 $\text{left_A}$ 和 $\text{left_B}$ 放入一个集合，并将 $\text{right_A}$ 和 $\text{right_B}$ 放入另一个集合。再把这两个新的集合分别命名为 $\text{left_part}$ 和 $\text{right_part}$ ：

          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]

当 $\text{A}$ 和 $\text{B}$ 的总长度是偶数时，如果可以确认：
$\text{len}(\text{left_part}) = \text{len}(\text{right_part})$
$\max(\text{left_part}) \leq \min(\text{right_part})$
那么， $\{\text{A}, \text{B}\}$ 中的所有元素已经被划分为相同长度的两个部分，且前一部分中的元素总是小于或等于后一部分中的元素。中位数就是前一部分的最大值和后一部分的最小值的平均值：
$\text{median} = \frac{\text{max}(\text{left}\_\text{part}) + \text{min}(\text{right}\_\text{part})}{2}$
当 $\text{A}$ 和 $\text{B}$ 的总长度是奇数时，如果可以确认：
$\text{len}(\text{left_part}) = \text{len}(\text{right_part})+1$
$\max(\text{left_part}) \leq \min(\text{right_part})$
那么， $\{\text{A}, \text{B}\}$ 中的所有元素已经被划分为两个部分，前一部分比后一部分多一个元素，且前一部分中的元素总是小于或等于后一部分中的元素。中位数就是前一部分的最大值：
$\text{median} = \text{max}(\text{left}\_\text{part})$
第一个条件对于总长度是偶数和奇数的情况有所不同，但是可以将两种情况合并。第二个条件对于总长度是偶数和奇数的情况是一样的。要确保这两个条件，只需要保证：

$i + j = m - i + n - j$ 当 $m+n$ 为偶数）或 $i + j = m - i + n - j + 1$ （当 $m+n$ 为奇数）。等号左侧为前一部分的元素个数，等号右侧为后一部分的元素个数。将 $i$ 和 $j$ 全部移到等号左侧，我们就可以得到 $i+j = \frac{m + n + 1}{2}$ 。这里的分数结果只保留整数部分。
，。如果我们规定的长度小于等于的长度，即。这样对于任意的，都有，那么我们在的范围内枚举并得到，就不需要额外的性质了。
- 如果 $\text{A}$ 的长度较大，那么我们只要交换 $\text{A}$ 和 $\text{B}$ 即可。
- 如果 $m > n$ ，那么得出的 $j$ 有可能是负数。
$\text{B}[j-1] \leq \text{A}[i]$ 以及 $\text{A}[i-1] \leq \text{B}[j]$ ，即前一部分的最大值小于等于后一部分的最小值。

为了简化分析，假设 $\text{A}[i-1], \text{B}[j-1], \text{A}[i], \text{B}[j]$ 总是存在。对于 $i=0$ 、 $i=m$ 、 $j=0$ 、 $j=n$ 这样的临界条件，我们只需要规定 $A[-1]=-\infty$ 、 $A[m]=\infty$ 、 $B[-1]=-\infty$ 、 $B[n]=\infty$ 即可。这也是比较直观的：当一个数组不出现在前一部分时，对应的值为负无穷，就不会对前一部分的最大值产生影响；当一个数组不出现在后一部分时，对应的值为正无穷，就不会对后一部分的最小值产生影响。
所以我们需要做的是：

在 $[0, m]$ 中找到 $i$ ，使得：
$\qquad \text{B}[j-1] \leq \text{A}[i]$ 且 $\text{A}[i-1] \leq \text{B}[j]$ ，其中 $j = \frac{m + n + 1}{2} - i$

我们证明它等价于：

在 $[0, m]$ 中找到最大的 $i$ ，使得：
$\qquad \text{A}[i-1] \leq \text{B}[j]$ ，其中 $j = \frac{m + n + 1}{2} - i$

这是因为：

当 $i$ 从 $0 \sim m$ 递增时， $\text{A}[i-1]$ 递增， $\text{B}[j]$ 递减，所以一定存在一个最大的 $i$ 满足 $A[i-1] \leq B[j]$ ；
如果 $i$ 是最大的，那么说明 $i+1$ 不满足。将 $i+1$ 带入可以得到 $A[i] > B[j-1]$ ，也就是 $B[j - 1] < A[i]$ ，就和我们进行等价变换前 $i$ 的性质一致了（甚至还要更强）。

因此我们可以对 $i$ 在 $[0, m]$ 的区间上进行二分搜索，找到最大的满足 $A[i-1] \leq B[j]$ 的 $i$ 值，就得到了划分的方法。此时，划分前一部分元素中的最大值，以及划分后一部分元素中的最小值，才可能作为就是这两个数组的中位数。

def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
    if len(nums1) > len(nums2):
        return self.findMedianSortedArrays(nums2, nums1)

    infinty = 2**40
    m, n = len(nums1), len(nums2)
    left, right, ansi = 0, m, -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) // 2
        j = (m + n + 1) // 2 - i

        # nums_im1, nums_i, nums_jm1, nums_j 分别表示 nums1[i-1], nums1[i], nums2[j-1], nums2[j]
        nums_im1 = (-infinty if i == 0 else nums1[i - 1])
        nums_i = (infinty if i == m else nums1[i])
        nums_jm1 = (-infinty if j == 0 else nums2[j - 1])
        nums_j = (infinty if j == n else nums2[j])
        if nums_im1 <= nums_j:
            median1, median2 = max(nums_im1, nums_jm1), min(nums_i, nums_j)
            left = i + 1
        else:
            right = i - 1
    return (median1 + median2) / 2 if (m + n) % 2 == 0 else median1

上帝之手（9行~12行）

思路与划分数组法相似。
思路
举个例子:

a = [2,3,6,8]
b = [1,4,5,7,9].

如果二者合并然后排序，我们得到[1,2,3,4,5,6,7,8,9]。根据题意，时间复杂度为O(log(m+n))，真实的合并操作会超时。不过我们可以假装二者已经合并了，从合并的数组中数出前4个数，第五个数5就是中位数。

那么这4个数应该是怎样的？假如我们从a数组取了 i 个数，再从b数组取4-i个数，那么该例中有下列5种取法：

      剩余部分      剩余部分    取用部分的最后一个数     取用部分的最后一个数
i    a[i, ...]   b[4-i, ...]    from b                  from a

0    [2,3,6,8]           [9]    7                       None
1      [3,6,8]         [7,9]    5                       2
2        [6,8]       [5,7,9]    4                       3
3          [8]     [4,5,7,9]    1                       6
4           []   [1,4,5,7,9]    None                    8

为了找到中位数，以下两条必须符合：

b数组的取用部分的最后一个数，不能大于a数组剩余部分的数；
a数组的取用部分的最后一个数，不能大于b数组剩余部分的数；

前两个取法 (i=0和i=1) 违反了第一条规则。随着i增加, b数组的取用部分的最后一个数缩小，a数组剩余部分的第一个数增加。这是一个单调增加过程。于是我们总会得到某种取法，符合两条规则。实际上，如果i从小到大寻找，只要找到符合第一条规则的i就会符合第二天规则。我们可以用二分法查找。在代码中，我们使用a[i]表示a数组剩余部分的第一个数，b[after-i-1]表示b数组的取用部分的最后一个数，不断检查是否符合第一条规则。

找出了前4个数，我们剩下[6,8]和[5,7,9]。我把两个数组各取出前两个元素，并把这4个数排序。在例子中这4个数是[5,6,7,8]。如果ab两个数组共有奇数个元素，我们只需取第一个数5作为中位数。如果元素个数是偶数 (比如我们在第一个数组中增补一个元素10)，我们只需取前两个数求均值，比如 (5+6)/2.0。为了代码的简洁性，合并两种情况，在奇数情况下不直接取5，而是计算 (5+5)/2.0。

如果利用Python自带的bisect库函数，则只需9行。

def findMedianSortedArrays(self, nums1, nums2):
    a, b = sorted((nums1, nums2), key=len)
    m, n = len(a), len(b)
    after = (m + n - 1) // 2
    class Range:
        def __getitem__(self, i):
            return after-i-1 < 0 or a[i] >= b[after-i-1]
    i = bisect.bisect_left(Range(), True, 0, m)
    nextfew = sorted(a[i:i+2] + b[after-i:after-i+2])
    return (nextfew[0] + nextfew[1 - (m+n)%2]) / 2.0

如果不利用Python自带的bisect库函数，需要12行。

def findMedianSortedArrays(self, nums1, nums2):
    a, b = sorted((nums1, nums2), key=len)
    m, n = len(a), len(b)
    after = (m + n - 1) // 2
    lo, hi = 0, m
    while lo < hi:
        i = (lo + hi) // 2
        if after-i-1 < 0 or a[i] >= b[after-i-1]:
            hi = i - 1
        else:
            lo = i + 1
    nextfew = sorted(a[lo:lo+2] + b[after-lo:after-lo+2])
    return (nextfew[0] + nextfew[1 - (m+n)%2]) / 2.0

需要注意，after=(m+n-1)//2而不是(m+n+1)//2。

Stefan Pochmann 的上帝之手（6）寻找两个正序数组

寻找两个正序数组的中位数

一、二分查找（22行）

划分数组（19行）

上帝之手（9行~12行）

猜你喜欢

热点阅读