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

news2025/5/22 19:52:04

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

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

📌 题目描述

给定两个正序（从小到大）排列的数组 nums1 和 nums2，要求在 不合并数组 的前提下，找出它们的中位数，并且时间复杂度必须为 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/2 和 total/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