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

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

题目链接

O(m + n)算法

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int n = nums1.length + nums2.length;
        int mid1 = (n + 1) / 2;
        int mid2 = (n + 2) / 2;
        int p1 = 0, p2 = 0, num1 = 0, num2 = 0;
        for(int i = 1; i <= n; i++){
            int num;
            if(p1 < nums1.length && (p2 == nums2.length || nums1[p1] <= nums2[p2])){
                num = nums1[p1];
                p1++;
            }
            else{
                num = nums2[p2];
                p2++;
            }
            if(i == mid1){
                num1 = num;
            }
            if(i == mid2){
                num2 = num;
                break;
            }
        }    
        return (num1 + num2) * 1.0 / 2; 
    }
}

先根据两个数组的长度和，算出中位数位置，再用两个指针分别遍历两个数组，找出中位数位置的两个数，求平均值。过程中需要注意一下几点：

• 如果长度和为奇数，设置两个中位数相同即可
• 要注意判断两个指针是否到头
• 找到两个中位数后可以提早break（实测最终运行时间与更快的算法相差无几）
• 返回的是double，记得先做类型转换

O(log(m+n))算法

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int n = nums1.length + nums2.length;
        if(n % 2 == 1)
            return find(nums1, 0, nums2, 0, (n >> 1) + 1);
        return (find(nums1, 0, nums2, 0, n >> 1) + find(nums1, 0, nums2, 0, (n >> 1) + 1)) * 1.0 / 2;
    }
    public int find(int[] nums1, int i, int[] nums2, int j, int k){
        if(i >= nums1.length){
            return nums2[j + k - 1];
        }
        if(j >= nums2.length){
            return nums1[i + k - 1];
        }
        if(k == 1){
            return Math.min(nums1[i], nums2[j]);
        }
        int mid1 = i + (k >> 1) - 1, mid2 = j + (k >> 1) - 1;
        int midNum1 = (mid1 < nums1.length) ? nums1[mid1] : Integer.MAX_VALUE;
        int midNum2 = (mid2 < nums2.length) ? nums2[mid2] : Integer.MAX_VALUE;
        if(midNum1 < midNum2){
            return find(nums1, mid1 + 1, nums2, j, k - (k >> 1));
        }
        return find(nums1, i, nums2, mid2 + 1, k - (k >> 1));
    }
}

算法要复杂一些，但实测没有快很多，可能是数据中m和n太小，而递归的耗时和常数系数太大。

具体做法就是，如果我们要找第k个数，可以先分别找出两个数组中的第k/2个数，比较大小。例如nums1中的第k/2个数较小，那么答案肯定不在num1的前k/2个数里，这样直接排除一半。但是比起O(m + n)的算法来说，多了每次比较的时间，而且总数(m+n)/2是小于等于1000的，所以效果并不明显。而且总数为奇数的时候，要找两次，拖慢速度。

当然算法中还要注意边界情况，以及剪枝。

此外，写了个非递归的版本。

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int n = nums1.length + nums2.length;
        if(n % 2 == 1)
            return find(nums1, nums2, (n >> 1) + 1);
        return (find(nums1, nums2, n >> 1) + find(nums1, nums2, (n >> 1) + 1)) * 1.0 / 2;
    }
    public int find(int[] nums1, int[] nums2, int k){
        int i = 0, j = 0;
        while(true){
            if(i >= nums1.length){
                return nums2[j + k - 1];
            }
            if(j >= nums2.length){
                return nums1[i + k - 1];
            }
            if(k == 1){
                return Math.min(nums1[i], nums2[j]);
            }
            int mid1 = i + (k >> 1) - 1, mid2 = j + (k >> 1) - 1;
            int midNum1 = (mid1 < nums1.length) ? nums1[mid1] : Integer.MAX_VALUE;
            int midNum2 = (mid2 < nums2.length) ? nums2[mid2] : Integer.MAX_VALUE;
            if(midNum1 < midNum2){
                i = mid1 + 1;
                k -= (k >> 1);
            }
            else{
                j = mid2 + 1;
                k -= (k >> 1);
            }
        }
    }
}

改进空间

这个算法最大的问题就是，m+n为奇数时，要找两次。改进的方法很简单，因为找两次时，找的是k和k+1，所以在第一次的基础上去找就行了。但是要额外判断很多东西，懒得写了。