[LeetCode]4. 寻找两个正序数组的中位数

题目

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

解题思路

中位数的作用是将一个集合划分为两个长度相等的子集，其中一个子集中的元素总是大于另一个子集中的元素。因此，可以在任意位置 i,j 将A和B分成两个部分：
A数组：

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

B数组：

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

将 left_A 和 left_B 放入一个集合，并将 right_A 和 right_B 放入另一个集合。 再把这两个新的集合分别命名为 left_part 和 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]

当 A 和 B 的总长度是偶数时，如果有：

1）len(left_part) = len(right_part)
2）max(left_part) ≤ min(right_part)

那么 {A,B} 中的所有元素已经被划分为相同长度的两个部分，且前一部分中的元素总是小于或等于后一部分中的元素。中位数就是前一部分的最大值和后一部分的最小值的平均值：

median = {max(left_part)+min(right_part)} / 2

当 A 和 B 的总长度是奇数时，如果有：

1）len(left_part) = len(right_part)+1
2）max(left_part) ≤ min(right_part)

那么 {A,B} 中的所有元素已经被划分为两个部分，前一部分比后一部分多一个元素，且前一部分中的元素总是小于或等于后一部分中的元素。中位数就是前一部分的最大值：

median = max(left_part)

要确保以上两个条件成立，只需要保证：
1）i+j = m-i+n-j（当 m+n 为偶数）或 i+j = m-i+n-j+1（当 m+n 为奇数）。等号左侧为前一部分的元素个数，等号右侧为后一部分的元素个数。将 i 和 j 全部移到等号左侧，我们就可以得到 i+j = (m+n+1)/2，此分数形式对于奇数或偶数的情况都成立。
2）0≤i≤m，0≤j≤n。如果我们规定 A 的长度小于等于 B 的长度，即 m≤n。这样对于任意的 i ∈ [0,m]，都有 j = (m+n+1)/2−i ∈ [0,n]，那么我们在 [0,m] 的范围内枚举 i 并得到 j，就不需要额外的性质了。如果出现 A 的长度大于 B 的长度，即 m>n的情况，那么得出的 j 有可能是负数，那么只要交换A 和 B 即可，保证 A 的长度小于等于 B 的长度。
3）B[j−1]≤A[i] 以及 A[i−1]≤B[j]，即前一部分的最大值小于等于后一部分的最小值。

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

所以我们需要做的是：

在 [0,m] 中找到 i，使得：B[j−1]≤A[i] 且 A[i−1]≤B[j]，其中 j = (m+n+1)/2−i

等价于：

在 [0,m] 中找到最大的 i，使得：A[i−1]≤B[j]，其中 j = (m+n+1)/2−i

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

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

复杂度分析：
时间复杂度：O(logmin(m,n)))，其中 m 和 n 分别是数组 nums1 和nums2 的长度。查找的区间是[0,m]，而该区间的长度在每次循环之后都会减少为原来的一半。所以，只需要执行 logm 次循环。由于每次循环中的操作次数是常数，所以时间复杂度为 O(logm)。由于我们可能需要交换 nums1 和 nums2 使得 m≤n，因此时间复杂度是O(logmin(m,n)))。
空间复杂度：O(1)。

代码

Python代码如下：

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m = len(nums1)
        n = len(nums2)
        # 保证 m<=n
        if m>n:
            return self.findMedianSortedArrays(nums2, nums1)
        # 分割线左边的所有元素需要满足的个数 m + (n - m + 1) / 2;
        totalLeft = (m+n+1)//2
        # 在 nums1 的区间 [0, m] 里查找恰当的分割线，
        # 使得 nums1[i - 1] <= nums2[j] && nums2[j - 1] <= nums1[i]。
        left = 0
        right = m
        while left<right:
            # 前一部分包含 nums1[0 .. i-1] 和 nums2[0 .. j-1]
            # 后一部分包含 nums1[i .. m-1] 和 nums2[j .. n-1]
            i = (left+right+1)//2
            j = totalLeft - i
            if nums1[i-1] > nums2[j]:
                #下一轮搜索的区间 [left, i - 1]
                right = i-1
            else:
                # 下一轮搜索的区间 [i, right]
                left = i
        i = left
        j = totalLeft - i

        nums1LeftMax = (-float('inf') if i==0 else nums1[i-1])
        nums1RightMin = (float('inf') if i==m else nums1[i])
        nums2LeftMax = (-float('inf') if j==0 else nums2[j-1])
        nums2RightMin = (float('inf') if j==n else nums2[j])

        # 总长度为奇数
        if (m+n)%2 == 1:
            return max(nums1LeftMax, nums2LeftMax)
        # 总长度为偶数
        else:
            return (max(nums1LeftMax, nums2LeftMax)+min(nums1RightMin, nums2RightMin))/2

Java代码如下：

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int m = nums1.length;
        int n = nums2.length;
        if(m>n){
            return findMedianSortedArrays(nums2, nums1);
        }
        int totalLeft = (m+n+1)/2;
        int left = 0, right = m;
        while(left<right){
            int i = (left+right+1)/2;
            int j = totalLeft - i;
            if(nums1[i-1]>nums2[j]){
                right = i-1;
            }else{
                left = i;
            }
        }
        int i = left;
        int j = totalLeft - i;
        
        int nums1LeftMax = i == 0 ? Integer.MIN_VALUE : nums1[i-1];
        int nums1RightMin = i == m ? Integer.MAX_VALUE : nums1[i];
        int nums2LeftMax = j == 0 ? Integer.MIN_VALUE : nums2[j-1];
        int nums2RightMin = j == n ? Integer.MAX_VALUE : nums2[j];

        if((m+n)%2 == 1){
            return Math.max(nums1LeftMax, nums2LeftMax);
        }else{
            return (double)(Math.max(nums1LeftMax, nums2LeftMax) + Math.min(nums1RightMin, nums2RightMin)) / 2;
        }
    }
}