每日算法 -【Swift 算法】寻找两个有序数组的中位数(O(log(m+n)))详细讲解版

用 Swift 寻找两个有序数组的中位数(O(log(m+n)))详细讲解版

寻找两个有序数组的中位数,是 LeetCode 上非常经典的一道题,难度为 困难(Hard),但它的本质是一个 二分查找 的变形应用。

题目描述

给定两个正序(从小到大)排列的数组 nums1nums2,要求在 不合并数组 的前提下,找出它们的中位数,并且时间复杂度必须为 O(log(m+n))

示例 1:

输入: nums1 = [1, 3], nums2 = [2]
合并后:[1, 2, 3]
中位数:2

示例 2:

输入: nums1 = [1, 2], nums2 = [3, 4]
合并后:[1, 2, 3, 4]
中位数:(2 + 3)/2 = 2.5

⚠️ 暴力解法(不推荐)

直接合并两个数组然后排序,再取中位数,这种做法时间复杂度是 O(m+n),不满足题目要求。

func findMedianSortedArraysBruteForce(_ nums1: [Int], _ nums2: [Int]) -> Double {
    // 1. 合并两个数组
    let merged = (nums1 + nums2).sorted()
    
    // 2. 获取总长度
    let count = merged.count
    
    // 3. 判断奇偶并返回中位数
    if count % 2 == 1 {
        // 奇数,直接返回中间的数
        return Double(merged[count / 2])
    } else {
        // 偶数,返回中间两个数的平均值
        let mid1 = merged[count / 2 - 1]
        let mid2 = merged[count / 2]
        return Double(mid1 + mid2) / 2.0
    }
}

✅ 二分查找解法(推荐)

我们要寻找的是第 k 小的数,而中位数正是数组中间的那个数。所以我们把问题转化为:

“在两个有序数组中找第 k 小的数”。


中位数转化为第 k 小数

设两个数组总长度为 total = m + n

  • 如果 total 是奇数,返回第 (total+1)/2 小的数;
  • 如果 total 是偶数,返回第 total/2total/2 + 1 小的数的平均值。

Swift 实现(附详细注释)

func findMedianSortedArrays(_ nums1: [Int], _ nums2: [Int]) -> Double {
    
    // 获取第 k 小的元素
    func getKthElement(_ k: Int) -> Int {
        var index1 = 0, index2 = 0
        var k = k
        
        let m = nums1.count, n = nums2.count

        while true {
            // 如果 nums1 已经全部用完了,直接从 nums2 取
            if index1 == m {
                return nums2[index2 + k - 1]
            }
            // 如果 nums2 已经全部用完了,直接从 nums1 取
            if index2 == n {
                return nums1[index1 + k - 1]
            }
            // 如果 k == 1,说明我们只需要找最小的那一个
            if k == 1 {
                return min(nums1[index1], nums2[index2])
            }

            // 正常情况下,从两个数组中各取 k/2 个元素进行比较
            let newIndex1 = min(index1 + k / 2 - 1, m - 1)
            let newIndex2 = min(index2 + k / 2 - 1, n - 1)

            let pivot1 = nums1[newIndex1]
            let pivot2 = nums2[newIndex2]

            if pivot1 <= pivot2 {
                // 舍弃 nums1[index1...newIndex1]
                k -= (newIndex1 - index1 + 1)
                index1 = newIndex1 + 1
            } else {
                // 舍弃 nums2[index2...newIndex2]
                k -= (newIndex2 - index2 + 1)
                index2 = newIndex2 + 1
            }
        }
    }

    // 主逻辑
    let totalLength = nums1.count + nums2.count

    if totalLength % 2 == 1 {
        // 奇数时:返回中间那个数
        return Double(getKthElement((totalLength + 1) / 2))
    } else {
        // 偶数时:返回中间两个数的平均值
        let mid1 = getKthElement(totalLength / 2)
        let mid2 = getKthElement(totalLength / 2 + 1)
        return Double(mid1 + mid2) / 2.0
    }
}

测试代码

print(findMedianSortedArrays([1, 3], [2]))       // 输出: 2.0
print(findMedianSortedArrays([1, 2], [3, 4]))    // 输出: 2.5
print(findMedianSortedArrays([0, 0], [0, 0]))    // 输出: 0.0
print(findMedianSortedArrays([], [1]))           // 输出: 1.0
print(findMedianSortedArrays([2], []))           // 输出: 2.0

时间复杂度分析

  • 每次比较会排除 k/2 个元素,时间复杂度为 O(log k)
  • 因为 k 最多是 m+n,所以总复杂度是 O(log(m+n)),符合要求。

✅ 总结

特性 描述
输入数组 有序
允许合并吗 ❌ 不允许
解法 二分查找
时间复杂度 O(log(m+n))
空间复杂度 O(1)

这种算法很有代表性,它说明了:二分查找不仅仅适用于一个数组,还可以延伸到两个数组,只要我们学会“不断舍弃不可能的部分”。


推荐练习

  • LeetCode 4. Median of Two Sorted Arrays
  • 寻找两个有序数组的第 K 小数

如果你觉得这篇文章对你有帮助,欢迎 收藏 ⭐,评论交流 !

你可能感兴趣的:(算法,算法,swift,开发语言)