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

· 约 5 分钟阅读 · 次阅读 leetcode

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

算法的时间复杂度应该为 O(log (m+n)) 。

示例 1:

输入:nums1 = [1,3], nums2 = [2]
输出:2.00000
解释:合并数组 = [1,2,3] ,中位数 2

示例 2:

输入:nums1 = [1,2], nums2 = [3,4]
输出:2.50000
解释:合并数组 = [1,2,3,4] ,中位数 (2 + 3) / 2 = 2.5

双指针法直接查找中位数

正常的解法是将两个数组合并,然后得到中位数,但是其实不用存储新数组,可以直接用双指针分别遍历, 分别取两个数组中较小的数,累积遍历(m+n)//2+1次即可。 考虑到奇数和偶数的差异,需要记录两个结果值。这里用left和right记录。 同时,还需要考虑如果有一个数组遍历完的情况

为了注重代码可读性,调整代码如下:

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        count = 0
        i = 0
        j = 0
        m = len(nums1)
        n = len(nums2)
        left = -1
        right = -1
        while count <= (m+n)//2 and i < m and j < n:
            left = right
            if nums1[i] < nums2[j]:
                right = nums1[i]
                i += 1
            else:
                right = nums2[j]
                j += 1
            count += 1
        
        while count <= (m+n)//2 and i < m:
            left = right
            right = nums1[i]
            i += 1
            count += 1
        while count <= (m+n)//2 and j < n:
            left = right
            right = nums2[j]
            j += 1
            count += 1
            # if i < m and (j >= n or nums1[i] < nums2[j]):
            #     right = nums1[i]
            #     i += 1
            # else:
            #     right = nums2[j]
            #     j += 1
            # count += 1
        if (m+n) % 2 == 0:
            ans = (left + right) / 2
        else:
            ans = right
        return ans

二分查找法-划分数组

时间复杂度:O(log min(m,n)) 空间复杂度:O(1)

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int m = nums1.size();
        int n = nums2.size();
        // 关键点1:让m始终比n小,可以使得我们以第一个数组为基准二分查找时,第二个数组分割线两侧一定有元素,也就是下标j不会越界
        // 特殊情况:
        // 1. 数组数量不等,A数组左边无元素
        // 2. 数组数量不等,A数组右边无元素
        // 3. 数组数量相等,A数组左边无元素,B数组右边无元素
        // 4. 数组数量相等,A数组右边无元素,B数组左边无元素
        if (m > n){
            return findMedianSortedArrays(nums2, nums1);
        }
        // 关键点2: 左半部分元素的个数
        // 如果是偶数,向下取整,正好一半;如果是奇数,让左边比右边多一个元素
        // 二分查找的时候,可以只调整第一个数组中的分割线位置,然后根据总长度控制第二个数组中的分割线位置
        int leftPartTotal = (m+n+1)/2;

        int left = 0;
        int right = m;
        int i = 0;
        int j = 0;

        // 第一个数组至少有1个元素,才会进入while循环
        // left == right时退出循环
        // 初始化循环, 也就是i的范围:[0, m]
        while (left < right){
            // 关键点3
            // i = (left + right + 1)/2 为了避免两个整数相加溢出,改为起点+增量计算中点的方式
            // i是第一个数组中分割线右边的第一个元素
            // j是第二个数组中分割线右边的第一个元素
            // 即,i,j所在的位置分别是分割线右边的第一个元素
            int i = left + (right-left+1)/2;
            int j = leftPartTotal - i;
            if (nums1[i-1] > nums2[j]){
                // 下一轮循环[left, i-1]
                right = i - 1;
            }else{
                // 下一轮循环[i, right]
                // 考虑只有两个元素的时候,需要让上面取中位数的地方上取整,避免进入死循环
                left = i;
            }
        }

        i = left;
        j = leftPartTotal - i;
        int firstLeft =  i == 0 ? INT_MIN : nums1[i-1];
        int firstRight = i == m ? INT_MAX : nums1[i];
        int secondLeft = j == 0 ? INT_MIN : nums2[j-1];
        int secondRight = j== n ? INT_MAX : nums2[j];     

        int leftMax = max(firstLeft, secondLeft);
        if ((m+n) % 2 == 1){
            return leftMax;
        }else{
            int rightMin = min(firstRight, secondRight);
            return float(leftMax + rightMin)/2;
        }

    }
};

python3重写一次

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m = len(nums1)
        n = len(nums2)
        if m > n:
            return self.findMedianSortedArrays(nums2, nums1)
        
        total_left = (m+n+1)//2
        left = 0
        right = m
        while left < right:
            i = left + (right - left + 1) // 2
            j = total_left - i
            # i的范围是[1, m], 所以i-1不会越界
            if nums1[i-1] > nums2[j]:
                # A数组左边大,说明分割线太靠右了,需要左移
                right = i - 1
            else:
                left = i
        
        i = left
        j = total_left - i
        first_left = nums1[i-1] if i > 0 else float('-inf')
        first_right = nums1[i] if i < m else float('inf')
        second_left = nums2[j-1] if j > 0 else float('-inf')
        second_right = nums2[j] if j < n else float('inf')

        left_max = max(first_left, second_left)
        right_min = min(first_right, second_right)
        return left_max if (m+n)%2==1 else (left_max+right_min)/2

二分查找法-查找第k大的数

时间复杂度O(log(m+n))

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m = len(nums1)
        n = len(nums2)
        if (m+n) % 2 == 0:
            # 偶数个数的时候,需要找两个数,然后求平均
            target1 = self.getKthElement(nums1,nums2,(m+n) // 2)
            target2 = self.getKthElement(nums1,nums2,(m+n) // 2+1)
            return (target1+target2)/2
        else:
            return self.getKthElement(nums1,nums2,(m+n)//2 + 1)
        
    
    def getKthElement(self, nums1, nums2, k):
        m = len(nums1)
        n = len(nums2)
        index1 = 0
        index2 = 0
        while True:
            if index1 == m:
                # k已经发生改变,不能使用m-k+1做索引
                return nums2[index2+k-1]
            if index2 == n:
                return nums1[index1+k-1]
            if k == 1:
                return min(nums1[index1], nums2[index2])

            new_index1 = min(index1 + k//2 - 1, m-1)
            new_index2 = min(index2 + k//2 - 1, n-1)
            pivot1 = nums1[new_index1]
            pivot2 = nums2[new_index2]
            # 每次排除k//2 - 1个数,直到寻找最后一个元素
            if pivot1 < pivot2:
                k = k - (new_index1 - index1 + 1)
                index1 = new_index1 + 1
            else:
                k = k - (new_index2 - index2 + 1)
                index2 = new_index2 + 1