python-第k小的数和嵌套函数

记录一下，用以自学

题目

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

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/median-of-two-sorted-arrays
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

 

为了达到时间复杂度O(log(m+n)的要求，就需要用二分法来做查询遍历。

总体思想很简单

• 首先根据两个数组的长度找到中位数是第k小个数，然后去每个数组numsi中的[k//2]个数，如果【k//2】超过数组的长度就去最后一个数（记作ni）。
• 然后比较nums1[n1]和nums2[n2],由于数组中本来就是正序的，所以小的那个数最大也是第k-1小，所以可以舍弃。
• 重复过程知道找到k-1个数

但是在代码实现中，我遇到一个麻烦，如何处理两个数组长度是偶数的情况。脑子短路没想到求两次，求最小k-1，和最小k然后求平均值。这个时候嵌套函数就来了，比起再写一个函数，嵌套函数更加简洁和清晰。

最后成功的代码：

from typing import List
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        def findKthElement(k):#找到两个数组合并后的第k个元素
            k=int(k)
            index1,index2=0,0  # nums1和nums2的可选择的起始下标
            m1,m2=len(nums1),len(nums2)
            while(1):
                if index1==m1:
                    return nums2[index2+k-1]
                elif index2==m2:
                    return nums1[index1+k-1]
                if(k==1):
                    return min(nums1[index1],nums2[index2])
                else:
                    # 一般情况
                    newIndex1=min(index1+k//2-1,m1-1)
                    newIndex2=min(index2+k//2-1,m2-1)
                    p1,p2=nums1[newIndex1],nums2[newIndex2]
                    if(p1<=p2):
                        k=k-(newIndex1-index1+1)
                        index1=newIndex1+1
                    else:
                        k = k - (newIndex2 - index2 + 1)
                        index2 = newIndex2+1
        l1,l2=len(nums1),len(nums2)
        l=l1+l2
        if(l%2==0):
            return (findKthElement(l/2)+findKthElement(l/2+1))/2
        else:
            return findKthElement((l+1)//2)

 

注意：k==1的情况要单独拿到一般情况之外讨论，因为k==1时，

newIndex1=min(index1+k//2-1,m1-1)

会出现倒退到已经被排除的序号上，即 newIndex1=index1-1

 

对于这道题，后面有看了官方给的第二种解法划分数组

官方代码如下：

from typing import List
class Solution:
    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 = 0, m
        # 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
            print(i,j)
            # 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

解法的大概思想就是划分两个数组A成A1，A2，B成B1，B2，使得A的前一部分和B的前一部分属于小二分之一，A的后一部分和B的后一部分属于大二分之一

一开始我一直不明白循环的判断条件只是保证了A的前一部分小于B的后一部分，但是没有保证B的前半部分小于A的前半部分，如何实现划分后[A1,B1]全部小于[A2,B2]

实际上他是对A进行二分划分，每次根据判断结果去调整A划分的位置，使得A1恰好全部属于小二分之一，A2恰好全部属于大二分之一，这样根据个数划分成的B1，B2就会满足条件

即，在划分A的过程中，如果满足条件

nums_im1 <= nums_j

，left向左继续探索，知道不满足条件，所以保证了精确划分