2 二分查找 leetcode4 寻找两个有序数组的中位数

题目描述

给定两个大小为 mn 的有序数组 nums1nums2。请你找出这两个有序数组的中位数,并且要求算法的时间复杂度为 O(log(m + n))
你可以假设 nums1nums2 不会同时为空。

示例 1:

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

则中位数是 2.0

示例2:

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

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

官方难度

Hard

解决思路

最简单粗暴的解决思路,从左到右遍历两个数组,找到中位数的位置:

  1. 计算nums1nums2的中位数,mid=(len(nums1)+len(nums2))/2,如果为长度和为偶数,则为mid=(len(nums1)+len(nums2))/2mid-1
  2. 同时从左到右遍历两个数组,如果nums1>nums2, nums2的下标++,否则nums1的下标++,直到找到第mid个数。

根据上面的分析,我们可以很容易的得到直接方案流程如下:

  1. 初始化start1start2代表当前nums1nums2的下标;
  2. 遍历mid+1次,每次找第i大的数,i从0开始,最终要找到第mid或者mid+1的数。
  3. 注意nums1或者nums2的下标越界问题。
  4. 如果nums1[start1]>nums2[start2]说明第i大的数字为nums2[start2]
  5. 如果i>=mid,可以结束循环,根据两个数组长度总和计算返回的中位数。
    基于这个流程,我们可以实现类似下面的代码:
class Solution(object):
    def findMedianSortedArrays(self, nums1, nums2):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float
        """
        n_len = len(nums1) + len(nums2)
        mid = n_len / 2
        start1 = 0
        start2 = 0
        pre = 0
        current_val = 0
        #当前第[0,i]大的数,最终要求第mid和mid+1位置的数字
        for i in range(mid+1):
            pre = current_val
            if start1 > len(nums1)-1:
                current_val = nums2[start2]
                start2 += 1

            elif start2 > len(nums2)-1:   
                current_val = nums1[start1]
                start1 += 1

            elif nums1[start1] < nums2[start2]:
                current_val = nums1[start1]
                start1 += 1
                
            else:
                current_val = nums2[start2]
                start2 += 1
            if i >= mid:
                if n_len % 2 != 0: 
                    return current_val
                return float(current_val+pre)/2

优化

看到有序数组,马上想到二分查找的思路,假设中位数是求第mid个数,我们假设在第一个数组中求mid/2大的数字,第二个数组中求mid/2大的数字,然后比较这两个数字,如果数组1的结果大,其实可以保证数组2的第mid/2个数肯定是要小于两个数组合并后的第mid个数,因此可以排除第二个数组的前mid/2个数字,以输入nums1:[1,2,5,6], nums2:[3,5,6,7,9]为例子,
1.要求mid=(4+5+1)/2=5个数字,

  1. nums1中求第5/2大的数字,从nums2中求第5/2大的数字,
  2. nums1[2-1] = 2, nums2[2-1] = 5
  3. 可以将nums1中的1,2两个数字都去掉,也就变成求整个数组的[5,6], [3,5,6,7,9]第3大的数字。
  4. 一直求到第1大的数字时候,比较nums1[0]和nums2[0]的最小值则是最终结果。
  5. 需要注意数组越界,数组为空,下标和mid对应以及偶数长度数组和奇数长度数组的问题。

基于以上思路我们可以梳理流程如下:

  1. 计算nnums1的长度,mnums2的长度。
  2. 解决偶数长度和奇数长度的问题,通过求(n+m+1)/2和(n+m+2)/2两个位置的数字,将这两类问题合并为一类问题,最终结果为[(n+m+1)/2 位置的数字+(n+m+2)/2位置的数字]*0.5
  3. 定义递归函数,用于求第K大的数字
  4. 递归函数定义结束条件,k==1
  5. 判断是否数组越界
  6. 比较大小,然后递归求k - 去除数组的数量

基于这个流程,我们可以实现类似下面的代码:

class Solution(object):
    def findMedianSortedArrays(self, nums1, nums2):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float
        """
        n = len(nums1)
        m = len(nums2)
        if not nums1:
            return (nums2[(n+m+1)/2-1] + nums2[(n+m+2)/2-1])*0.5
        if not nums2:
            return (nums1[(n+m+1)/2-1] + nums1[(n+m+2)/2-1])*0.5
        def getKth(nums1, start1, end1, nums2, start2, end2, k):  
            """
            start1: 数组1第一个元素开始位置
            end1: 数组1最后一个元素位置
            start2: 数组2第一个元素开始位置
            end2: 数组2最后一个元素位置
            k:需要求的第K大元素
            """
            if end1 < start1: 
                return nums2[start2+k-1]
            if end2 < start2: 
                return nums1[start1+k-1]
            if k == 1: 
                return min(nums1[start1], nums2[start2])
    
            start1_temp = min(start1+k/2-1, len(nums1)-1)
            start2_temp = min(start2+k/2-1, len(nums2)-1)

            if nums2[start2_temp] > nums1[start1_temp]:
                return getKth(nums1, start1_temp+1, end1, nums2, start2, end2, k-(start1_temp - start1+1))
            else:
                return getKth(nums1, start1, end1, nums2, start2_temp+1, end2, k-(start2_temp - start2+1))
        return (getKth(nums1, 0, n-1, nums2, 0,  m-1, (n+m+1)/2) + getKth(nums1, 0, n-1, nums2, 0,  m-1, (n+m+2)/2))*0.5
        

总结

这题真是费了太大力气才调通了,主要原因是数组的下标太多了,一定要注意定义好每一个变量的循环不变量,另外这道题求中位数,更可怕的是有两种情况,这次直接通过(n+m+1)/2和(n+m+2)/2来完美融合两种情况。

  1. 合并奇偶问题: (n+m+1)/2和(n+m+2)/2
  2. 下标问题,明确循环不变量,end代表最后一个元素位置,start代表第一个元素位置,k代表第k大的数字位置。
  3. 越界问题:取数组下标和数组长度的最小值。
  4. 删除数组问题:修改start位置,同时更改k的大小。
  5. 空数组问题,递归函数要判断是否end>start。

你可能感兴趣的:(2 二分查找 leetcode4 寻找两个有序数组的中位数)