二分查找题目：寻找两个正序数组的中位数

题目

标题和出处

标题：寻找两个正序数组的中位数

出处：4. 寻找两个正序数组的中位数

难度

8 级

题目描述

要求

给定两个大小分别为 $\text{m}$ 和 $\text{n}$ 的升序数组 $\text{nums1}$ 和 $\text{nums2}$，返回这两个升序数组的中位数。

要求时间复杂度是 $\text{O(log (m + n))}$。

示例

示例 1：

输入：$\text{nums1 = [1,3], nums2 = [2]}$
输出：$\text{2.00000}$
解释：合并数组是 $\text{[1,2,3]}$，中位数是 $\text{2}$。

示例 2：

输入：$\text{nums1 = [1,2], nums2 = [3,4]}$
输出：$\text{2.50000}$
解释：合并数组是 $\text{[1,2,3,4]}$，中位数是 $\frac{\text{2}+\text{3}}{\text{2}}=\text{2.5}$。

数据范围

• $\text{nums1.length}=\text{m}$
• $\text{nums2.length}=\text{n}$
• $\text{0}\le \text{m}\le \text{1000}$
• $\text{0}\le \text{n}\le \text{1000}$
• $\text{1}\le \text{m}+\text{n}\le \text{2000}$
• ${\text{-10}}^{\text{6}}\le \text{nums1[i], nums2[i]}\le {\text{10}}^{\text{6}}$

解法一

思路和算法

已知两个升序数组的长度分别是 $m$ 和 $n$。计算两个升序数组的中位数可以转换成找到两个升序数组的所有元素中的第 $k$ 小元素，其中 $0\le k<m+n$。用 $\text{total}=m+n$ 表示两个升序数组的长度之和。当 $\text{total}$ 是奇数时，$k=\frac{\text{total}-1}{2}$，第 $k$ 小元素即为中位数；当 $\text{total}$ 是偶数时，分别取 $k=\frac{\text{total}}{2}-1$ 和 $k=\frac{\text{total}}{2}$，两次第 $k$ 小元素的平均数即为中位数。因此，根据两个升序数组的长度之和是奇数或偶数，执行一次或两次寻找第 $k$ 小元素的操作，即可得到中位数。

由于题目要求时间复杂度是 $O(\log (m+n))$，因此要求每次寻找第 $k$ 小元素的操作的时间复杂度是 $O(\log (m+n))$。需要使用二分查找实现。

用 $k$ 表示目标值在剩余元素中的序号（$k$ 从 $0$ 开始，序号为 $k$ 表示剩余元素中有 $k$ 个元素小于等于目标值），用 ${\text{index}}_1$ 和 ${\text{index}}_2$ 分别表示数组 ${\text{nums}}_1$ 和 ${\text{nums}}_2$ 的首个剩余元素的下标，初始时 ${\text{index}}_1$ 和 ${\text{index}}_2$ 都等于 $0$。剩余元素表示可能是目标值的元素，查找过程中将不可能是目标值的元素排除。

每次查找时，分别考虑两个数组的剩余元素中最小的 $\lceil \frac{k}{2}\rceil$ 个元素，共考虑 $k+1$ 个元素（当 $k$ 是奇数时）或 $k$ 个元素（当 $k$ 是偶数时），这些元素在两个数组中的下标范围分别是 ${\text{nums}}_1$ 的下标范围 $[{\text{index}}_1,{\text{endIndex}}_1]$ 和 ${\text{nums}}_2$ 的下标范围 $[{\text{index}}_2,{\text{endIndex}}_2]$，其中 ${\text{endIndex}}_1={\text{index}}_1+\lfloor \frac{k-1}{2}\rfloor$，${\text{endIndex}}_2={\text{index}}_2+\lfloor \frac{k-1}{2}\rfloor$。考虑 ${\text{nums}}_1[{\text{endIndex}}_1]$ 和 ${\text{nums}}_2[{\text{endIndex}}_2]$，其中的较大值是第 $k$ 小元素（当 $k$ 是奇数时）或第 $k-1$ 小元素（当 $k$ 是偶数时），因此其中的较小值一定不是第 $k$ 小元素。对于较小值所在的数组，可以将较小值以及较小值前面的元素全部排除。

需要注意的是，${\text{endIndex}}_1$ 和 ${\text{endIndex}}_2$ 不能超出数组下标范围。如果一个数组的剩余元素个数少于 $\lceil \frac{k}{2}\rceil$，则该数组中考虑的元素是该数组中的全部剩余元素。因此有 ${\text{endIndex}}_1=\min ({\text{index}}_1+\lfloor \frac{k-1}{2}\rfloor ,m-1)$，${\text{endIndex}}_2=\min ({\text{index}}_2+\lfloor \frac{k-1}{2}\rfloor ,n-1)$。

由此可以根据三种情况分别做相应的处理，缩小查找范围。

• 如果 ${\text{nums}}_1[{\text{endIndex}}_1]<{\text{nums}}_2[{\text{endIndex}}_2]$，则将 ${\text{nums}}_1$ 的下标范围 $[{\text{index}}_1,{\text{endIndex}}_1]$ 中的元素全部排除，排除的元素个数是 ${\text{endIndex}}_1-{\text{index}}_1+1$。

• 如果 ${\text{nums}}_1[{\text{endIndex}}_1]>{\text{nums}}_2[{\text{endIndex}}_2]$，则将 ${\text{nums}}_2$ 的下标范围 $[{\text{index}}_2,{\text{endIndex}}_2]$ 中的元素全部排除，排除的元素个数是 ${\text{endIndex}}_2-{\text{index}}_2+1$。

• 如果 ${\text{nums}}_1[{\text{endIndex}}_1]={\text{nums}}_2[{\text{endIndex}}_2]$，则处理方式和 ${\text{nums}}_1[{\text{endIndex}}_1]<{\text{nums}}_2[{\text{endIndex}}_2]$ 相同。

每次查找之后，将 $k$ 的值减去排除的元素个数，并将排除元素的数组的相应下标更新为该数组首个剩余元素的下标，具体做法如下：如果排除的是 ${\text{nums}}_1$ 中的元素，则将 ${\text{index}}_1$ 更新为 ${\text{endIndex}}_1+1$；如果排除的是 ${\text{nums}}_2$ 中的元素，则将 ${\text{index}}_2$ 更新为 ${\text{endIndex}}_2+1$。

二分查找的条件是 ${\text{index}}_1<m$，${\text{index}}_2<n$ 和 $k>0$。如果三个条件之一不满足，则二分查找结束，得到目标值。

• 如果 ${\text{index}}_1=m$，则剩余元素都在 ${\text{nums}}_2$ 中，目标值是 ${\text{nums}}_2[{\text{index}}_2+k]$。

• 如果 ${\text{index}}_2=n$，则剩余元素都在 ${\text{nums}}_1$ 中，目标值是 ${\text{nums}}_1[{\text{index}}_1+k]$。

• 如果 $k=0$，则剩余元素中的最小元素是目标值，目标值是 $\min ({\text{nums}}_1[{\text{index}}_1],{\text{nums}}_2[{\text{index}}_2])$。

以下用一个例子说明该解法。

两个数组是 ${\text{nums}}_1=[1,2,3,4,5]$，${\text{nums}}_2=[1,2,3,4,5,6,7,8,9,10]$，两个数组的长度分别是 $m=5$，$n=10$，长度之和是 $15$，$k=7$。初始时，${\text{index}}_1=0$，${\text{index}}_2=0$。

1. 根据 ${\text{index}}_1=0$，${\text{index}}_2=0$ 和 $k=7$ 计算得到 ${\text{endIndex}}_1=3$，${\text{endIndex}}_2=3$。由于 ${\text{nums}}_1[3]={\text{nums}}_2[3]$，因此将 ${\text{nums}}_1$ 的下标范围 $[0,3]$ 排除，排除 $4$ 个元素，更新得到 $k=3$，${\text{index}}_1=4$。

2. 根据 ${\text{index}}_1=4$，${\text{index}}_2=0$ 和 $k=3$ 计算得到 ${\text{endIndex}}_1=4$，${\text{endIndex}}_2=1$。由于 ${\text{nums}}_1[4]>{\text{nums}}_2[1]$，因此将 ${\text{nums}}_2$ 的下标范围 $[0,1]$ 排除，排除 $2$ 个元素，更新得到 $k=1$，${\text{index}}_2=2$。

3. 根据 ${\text{index}}_1=4$，${\text{index}}_2=2$ 和 $k=1$ 计算得到 ${\text{endIndex}}_1=4$，${\text{endIndex}}_2=2$。由于 ${\text{nums}}_1[4]>{\text{nums}}_2[2]$，因此将 ${\text{nums}}_2$ 的下标范围 $[2,2]$ 排除，排除 $1$ 个元素，更新得到 $k=0$，${\text{index}}_2=3$。

4. 此时 $k=0$，二分查找结束，${\text{nums}}_1[4]$ 和 ${\text{nums}}_2[3]$ 中的较小值 $4$ 即为目标值。

代码

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int m = nums1.length, n = nums2.length;
        int total = m + n;
        if (total % 2 == 1) {
            int medianIndex = (total - 1) / 2;
            return findKthSmallest(medianIndex, nums1, nums2);
        } else {
            int medianIndex1 = total / 2 - 1, medianIndex2 = total / 2;
            return (findKthSmallest(medianIndex1, nums1, nums2) + findKthSmallest(medianIndex2, nums1, nums2)) / 2.0;
        }
    }

    public int findKthSmallest(int k, int[] nums1, int[] nums2) {
        int m = nums1.length, n = nums2.length;
        int index1 = 0, index2 = 0;
        while (index1 < m && index2 < n && k > 0) {
            int endIndex1 = Math.min(index1 + (k - 1) / 2, m - 1);
            int endIndex2 = Math.min(index2 + (k - 1) / 2, n - 1);
            int num1 = nums1[endIndex1], num2 = nums2[endIndex2];
            if (num1 <= num2) {
                k -= endIndex1 - index1 + 1;
                index1 = endIndex1 + 1;
            } else {
                k -= endIndex2 - index2 + 1;
                index2 = endIndex2 + 1;
            }
        }
        if (index1 == m) {
            return nums2[index2 + k];
        } else if (index2 == n) {
            return nums1[index1 + k];
        } else {
            return Math.min(nums1[index1], nums2[index2]);
        }
    }
}

复杂度分析

• 时间复杂度：$O(\log (m+n))$，其中 $m$ 和 $n$ 分别是数组 ${\text{nums}}_1$ 和 ${\text{nums}}_2$ 的长度。每次寻找第 $k$ 小元素时，$k$ 的初始值是 $m+n$ 的一半附近的整数，每次查找将 $k$ 的值减小一半，因此时间复杂度是 $O(\log (m+n))$。

• 空间复杂度：$O(1)$。

解法二

思路和算法

解法一的时间复杂度是 $O(\log (m+n))$，该时间复杂度已经很低，但是这道题还存在时间复杂度更低的解法。

为了找到中位数，需要在数组 ${\text{nums}}_1$ 和 ${\text{nums}}_2$ 中分别找到分割点 ${\text{cut}}_1$ 和 ${\text{cut}}_2$，将每个数组分割成两个部分。

• 数组 ${\text{nums}}_1$ 被分割成下标范围 $[0,{\text{cut}}_1-1]$ 和下标范围 $[{\text{cut}}_1,m-1]$ 两部分，左边部分的长度是 ${\text{cut}}_1$。

• 数组 ${\text{nums}}_2$ 被分割成下标范围 $[0,{\text{cut}}_2-1]$ 和下标范围 $[{\text{cut}}_2,n-1]$ 两部分，左边部分的长度是 ${\text{cut}}_2$。

其中，$0\le {\text{cut}}_1\le m$，$0\le {\text{cut}}_2\le n$，即每个数组分割成的两个部分中可以有一个部分为空。

假设 ${\text{nums}}_1[-1]={\text{nums}}_2[-1]=-\infty$，${\text{nums}}_1[m]={\text{nums}}_2[n]=+\infty$，分割应满足以下两个条件。

• 两个数组的左边部分的最大值小于等于两个数组的右边部分的最小值，$\max ({\text{nums}}_1[{\text{cut}}_1-1],{\text{nums}}_2[{\text{cut}}_2-1])\le \min ({\text{nums}}_1[{\text{cut}}_1],{\text{nums}}_2[{\text{cut}}_2])$。

• 两个数组的左边部分的长度之和为两个数组的长度之和的一半向上取整，${\text{cut}}_1+{\text{cut}}_2=\lceil \frac{m+n}{2}\rceil$。

将两个数组的左边部分统称为前半部分，将两个数组的右边部分统称为后半部分，则前半部分的最大值小于等于后半部分的最小值，前半部分的元素个数为两个数组的长度之和的一半向上取整。

用 $\text{total}=m+n$ 表示两个升序数组的长度之和，用 $\text{lowerSize}=\lceil \frac{\text{total}}{2}\rceil$ 表示前半部分的元素个数。当 $\text{total}$ 是奇数时，中位数是前半部分的最大值；当 $\text{total}$ 是偶数时，中位数是前半部分的最大值与后半部分的最小值的平均数。

由于已知 ${\text{cut}}_1+{\text{cut}}_2=\text{lowerSize}$，因此可以在 ${\text{nums}}_1$ 中寻找 ${\text{cut}}_1$，当 ${\text{cut}}_1$ 确定之后 ${\text{cut}}_2$ 也可以确定。

寻找 ${\text{cut}}_1$ 可以使用二分查找实现。由于两个数组都是升序数组，${\text{nums}}_1[{\text{cut}}_1-1]\le {\text{nums}}_1[{\text{cut}}_1]$ 和 ${\text{nums}}_2[{\text{cut}}_2-1]\le {\text{nums}}_2[{\text{cut}}_2]$ 都满足，因此只需要满足 ${\text{nums}}_1[{\text{cut}}_1-1]\le {\text{nums}}_2[{\text{cut}}_2]$ 和 ${\text{nums}}_2[{\text{cut}}_2-1]\le {\text{nums}}_1[{\text{cut}}_1]$ 即可。二分查找需要查找满足 ${\text{nums}}_1[{\text{cut}}_1-1]\le {\text{nums}}_2[{\text{cut}}_2]$ 的最大下标 ${\text{cut}}_1$。

用 $\text{low}$ 和 $\text{high}$ 分别表示二分查找的下标范围的下界和上界，初始时 $\text{low}=0$，$\text{high}=m$。每次查找时，取 ${\text{index}}_1$ 为 $\text{low}$ 和 $\text{high}$ 的平均数向上取整，并得到 ${\text{index}}_2=\text{lowerSize}-{\text{index}}_1$，比较 ${\text{nums}}_1[{\text{index}}_1-1]$ 和 ${\text{nums}}_2[{\text{index}}_2]$ 的大小关系，调整查找的下标范围。

• 如果 ${\text{nums}}_1[{\text{index}}_1-1]\le {\text{nums}}_2[{\text{index}}_2]$，则 ${\text{cut}}_1\ge {\text{index}}_1$，因此在下标范围 $[{\text{index}}_1,\text{high}]$ 中继续查找。

• 如果 ${\text{nums}}_1[{\text{index}}_1-1]>{\text{nums}}_2[{\text{index}}_2]$，则 ${\text{cut}}_1<{\text{index}}_1$，因此在下标范围 $[\text{low},{\text{index}}_1-1]$ 中继续查找。

当 $\text{low}=\text{high}$ 时，查找结束，此时 $\text{low}$ 即为 ${\text{cut}}_1$。

得到 ${\text{cut}}_1$ 之后即可得到 ${\text{cut}}_2$，${\text{nums}}_1[{\text{cut}}_1-1]$ 和 ${\text{nums}}_2[{\text{cut}}_2-1]$ 中的最大值是前半部分的最大值，${\text{nums}}_1[{\text{cut}}_1]$ 和 ${\text{nums}}_2[{\text{cut}}_2]$ 中的最小值是后半部分的最小值。根据前半部分的最大值和后半部分的最小值即可计算中位数。

• 当 $\text{total}$ 是奇数时，中位数是前半部分的最大值。

• 当 $\text{total}$ 是偶数时，中位数是前半部分的最大值与后半部分的最小值的平均数。

该解法的时间复杂度是 $O(\log m)$，优于解法一的 $O(\log (m+n))$。

实现方面，由于只需要在一个数组中二分查找，因此可以选择较短的数组二分查找，时间复杂度是 $O(\log \min (m,n))$。

以下用一个例子说明上述过程。

两个数组是 ${\text{nums}}_1=[1,2,3,4,5]$，${\text{nums}}_2=[1,2,3,4,5,6,7,8,9,10]$，两个数组的长度分别是 $m=5$，$n=10$，长度之和是 $15$，前半部分的元素个数是 $8$。初始时，$\text{low}=0$，$\text{high}=5$。

1. 根据 $\text{low}=0$ 和 $\text{high}=5$ 计算得到 ${\text{index}}_1=3$，${\text{index}}_2=5$。由于 ${\text{nums}}_1[2]\le {\text{nums}}_2[5]$，因此将 $\text{low}$ 更新为 $3$。

2. 根据 $\text{low}=3$ 和 $\text{high}=5$ 计算得到 ${\text{index}}_1=4$，${\text{index}}_2=4$。由于 ${\text{nums}}_1[3]\le {\text{nums}}_2[4]$，因此将 $\text{low}$ 更新为 $4$。

3. 根据 $\text{low}=4$ 和 $\text{high}=5$ 计算得到 ${\text{index}}_1=5$，${\text{index}}_2=3$。由于 ${\text{nums}}_1[4]>{\text{nums}}_2[3]$，因此将 $\text{high}$ 更新为 $4$。

4. 此时 $\text{low}=\text{high}$，二分查找结束。根据 $\text{low}=4$ 计算得到 ${\text{cut}}_1=4$，${\text{cut}}_2=4$，前半部分的最大值是 $4$，后半部分的最小值是 $5$。由于两个数组的长度之和是奇数，因此中位数是前半部分的最大值，中位数是 $4$。

代码

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        return nums1.length <= nums2.length ? findMedian(nums1, nums2) : findMedian(nums2, nums1);
    }

    public double findMedian(int[] shorter, int[] longer) {
        int length1 = shorter.length, length2 = longer.length;
        int total = length1 + length2;
        int lowerSize = (total + 1) / 2;
        int low = 0, high = length1;
        while (low < high) {
            int index1 = low + (high - low + 1) / 2;
            int index2 = lowerSize - index1;
            int left1 = shorter[index1 - 1];
            int right2 = longer[index2];
            if (left1 <= right2) {
                low = index1;
            } else {
                high = index1 - 1;
            }
        }
        int cut1 = low, cut2 = lowerSize - low;
        int lower1 = cut1 == 0 ? Integer.MIN_VALUE : shorter[cut1 - 1];
        int lower2 = cut2 == 0 ? Integer.MIN_VALUE : longer[cut2 - 1];
        int higher1 = cut1 == length1 ? Integer.MAX_VALUE : shorter[cut1];
        int higher2 = cut2 == length2 ? Integer.MAX_VALUE : longer[cut2];
        int lowerMax = Math.max(lower1, lower2), higherMin = Math.min(higher1, higher2);
        if (total % 2 == 1) {
            return lowerMax;
        } else {
            return (lowerMax + higherMin) / 2.0;
        }
    }
}

复杂度分析

• 时间复杂度：$O(\log \min (m,n))$，其中 $m$ 和 $n$ 分别是数组 ${\text{nums}}_1$ 和 ${\text{nums}}_2$ 的长度。在较短的数组中二分查找，范围是 $[0,\min (m,n)]$，二分查找的次数是 $O(\log \min (m,n))$，每次查找的时间是 $O(1)$，因此时间复杂度是 $O(\log \min (m,n))$。

• 空间复杂度：$O(1)$。