首页手记 LeetCode 寻找两个有序数组的中位数

LeetCode 寻找两个有序数组的中位数

标签：

Python

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

题目来源：https://leetcode-cn.com/problems/median-of-two-sorted-arrays/submissions/

题目

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

解题思路

使用方法：递归；
本题数组有序，要求中位数，但可将本题转换为求第 k 个最小数。若为奇数，则第 k 位所在的数就是中位数；若为偶数，则第 k 小的数与第 k + 1 小的数之和的一半就是中位数；
因为数组是有序的，两个数组比较求第 k 小的数，不需逐个比较，一个个剔除。可以考虑一半一半的排除，每次排除 k/2 个数字。
每次比较两个数组中第 k/2 位数（当 k 为奇数时，向下取整），当出现其中一个数组的值较小时，说明这个数组包括第 k/2 位数以及前面的数字都比第 k 位数小，直接剔除；此时 k 也要减去剔除数字的个数，递归执行这个步骤；
递归跳出的条件，当其中一个数组为空的时候，或者当 k 为 1 的时候；
当 k 为 1 的时候，这时求第 1 小的数，只要比较两个数组首位数字较小即是所求的中位数；
当数组为空，且 k 不等于 1 时，这个时候，取其中不为空数组的第 k 小的数。

例子 1：

nums1： [1, 3, 5]
nums2： [1, 2, 3, 4, 5, 6, 7, 8]

这个例子中，假设要求的是第 6 小的数字。比较 k/2 也就是第 3 个数。

nums1 第 3 个数字为 5，nums2 第 3 个数字为 3，5 > 3，所以将下面数组的 3 个数字剔除掉，变成比较 [1, 3, 5] 与 [4, 5, 6, 7, 8]，这个时候剔除 3 个数字，所以 k 需减去 3，此时 k=3。继续比较；

k/2 不能整除，向下取整，为 1，比较两个数组第一个数字，此时 1<4，剔除第一个数组的 1，剩下 [3, 5] 和 [4, 5, 6, 7, 8]，k 为 3-1=2，再次比较；

k/2 为 1，3<4，剔除第一个数组的 3，剩下 [5] 和 [4, 5, 6, 7, 8]，k 为 2-1=1。

此时 k=1，也就是求第 1 小的数字，这个时候比较剩下两个数组的首位数字，去较小的数字即为第 1 小的数字，也就是 4。

例子 2：

这里考虑数组长度小于 k/2 的情况：

nums1: [1, 2]
nums2: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

同样求第 6 小的数。

解题思路如 例子 1，先比较第 k/2 位数的大小，k/2=3。由于第一个数组长度小于 3，这时直接取最末尾的数字 2，而第二个数组取第 3 位数，也就是 3，这里 2<3，直接剔除第一个数组的两个数字，k 则变为 6-2=4。

由于此时第一个数组为空，所以这个时候直接取另外一个数组第 k 小，也就是第 4 小数，也就是 4，即是本来要求的两个数组中第 6 小的数字。

因为每次循环都会减少 k/2 个元素，所以时间复杂度是O(log(k))，而 k 在本题中就是 (m+n)/2，所以最终复杂度就是 O(log(m+n))。

代码实现

class Solution:
    def findMedianSortedArrays(self, nums1, nums2) -> float:
        m = len(nums1)
        n = len(nums2)
        # 这里默认为偶数情况，若为奇数，则计算两次相同的 k 值
        med_left = (m + n + 1) // 2
        med_right = (m + n + 2) // 2

        return (get_k_num(nums1, 0, m-1, nums2, 0, n-1, med_left) + get_k_num(nums1, 0, m-1, nums2, 0, n-1, med_right)) / 2

def get_k_num(nums1, head1, end1, nums2, head2, end2, k):
        len1 = end1 - head1 + 1
        len2 = end2 - head2 + 1
        # 递归跳出条件
        if (len1 == 0):
            return nums2[head2 + k - 1]
        if (len2 == 0):
            return nums1[head1 + k - 1]
        if (k == 1):
            return min(nums1[head1], nums2[head2])
        
        i = head1 + min(len1, k // 2) - 1
        j = head2 + min(len2, k // 2) - 1

        if (nums1[i] > nums2[j]):
            return get_k_num(nums1, head1, end1, nums2, j + 1, end2, k - (j - head2 + 1))
        
        else:
            return get_k_num(nums1, i + 1, end1, nums2, head2, end2, k - (i - head1 + 1))