LeetCode-4. 寻找两个正序数组的中位数

问题分析

对于寻找两个正序数组的中位数，可以进一步普适化成寻找两个正序数组中第 k 小的数。面对这类问题，可以使用分治的思路，将查找的时间复杂度降低到 $O(log(m+n))$

算法描述

分：在两个数组中找到第 k 小的数，可以分解成比较两个数组的第 $\frac{k}{2}$ 个元素

治：记两个数组nums1和nums2的起始下标为idx1和idx2，两个数组的第 $\frac{k}{2}$ 个元素的下标为t1和t2

• 若nums1[t1] <= nums2[t2]，说明nums1位于t1左区间的元素不可能是第 k 小的数
• 若nums1[t1] > nums2[t2]，说明nums2位于t1左区间的元素不可能是第 k 小的数

合并：排除了一些不可能的元素后，重新更新 $k$ 和数组的起始下标，缩小搜索范围

• 若nums1[t1] <= nums2[t2]，nums1的起始下标从t1 + 1开始，原问题变为找两个数组第k - (t1 - idx + 1)个元素
• 若nums1[t1] > nums2[t2]，nums2的起始下标从t2 + 1开始，原问题变为找两个数组第k - (t2 - idx + 1)个元素

程序代码

C++

class Solution {
private:
    double findKth(vector<int>& nums1, vector<int>& nums2, int k) {
        int m = nums1.size(), n = nums2.size();
        int idx1 = 0, idx2 = 0;
        while( true ) {
            if( idx1 == m )  return nums2[idx2 + k - 1];
            if( idx2 == n )  return nums1[idx1 + k - 1];
            if( k == 1 )  return min(nums1[idx1], nums2[idx2]);

            int t1 = min(idx1 + k / 2 - 1, m - 1);
            int t2 = min(idx2 + k / 2 - 1, n - 1);
            if( nums1[t1] <= nums2[t2] ) {
                k -= t1 - idx1 + 1;
                idx1 = t1 + 1;
            }
            else {
                k -= t2 - idx2 + 1;
                idx2 = t2 + 1;
            }
        }
    }

public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int m = nums1.size(), n = nums2.size();
        if( (m + n) & 1 ) {
            return findKth(nums1, nums2, (m + n + 1) / 2);
        }
        double t1 = findKth(nums1, nums2, (m + n) / 2);
        double t2 = findKth(nums1, nums2, (m + n) / 2 + 1);
        return (t1 + t2) / 2;
    }
};

Go

func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
    m, n := len(nums1), len(nums2)
    if (m + n) % 2 == 1 {
        return float64(findKth(nums1, nums2, (m + n + 1) / 2))
    }
    t1 := float64(findKth(nums1, nums2, (m + n) / 2))
    t2 := float64(findKth(nums1, nums2, (m + n) / 2 + 1))
    return (t1 + t2) / 2
}

func findKth(nums1 []int, nums2 []int, k int) int {
    m, n := len(nums1), len(nums2)
    idx1, idx2 := 0, 0
    for {
        if idx1 == m {
            return nums2[idx2 + k - 1]
        }
        if idx2 == n {
            return nums1[idx1 + k - 1]
        }
        if k == 1 {
            return min(nums1[idx1], nums2[idx2])
        }

        t1 := min(idx1 + k/2 - 1, m - 1);
        t2 := min(idx2 + k/2 - 1, n - 1);

        if nums1[t1] <= nums2[t2] {
            k -= t1 - idx1 + 1
            idx1 = t1 + 1
        } else {
            k -= t2 -idx2 + 1
            idx2 = t2 + 1
        }
    }
}

func min(x, y int) int {
    if x <= y {
        return x
    }
    return y
}