[Leetcode #4]Median of Two Sorted Arrays 计算两个有序数组的中位数

原题地址:https://leetcode.com/problems/median-of-two-sorted-arrays/

题目要求是:给定两个有序数组nums1[m]和nums2[n],计算它们的中位数,要求算法复杂度是O(log(m+n))。举例:

nums1 = [1, 3], nums2 = [2], 中位数是2.0

nums1 = [1, 2], nums2 = [3, 4], 中位数是(2 + 3) / 2 = 2.5


这题乍一看很简单嘛,把俩数组并成一个数组,计算中位数不就行了。但是题目要求的算法复杂度是O(log(m+n)),显然不达标。

要达到对数级的算法复杂度,必然会联系到二分查找。怎么二分查找呢?举个例子看一看:

nums1 = [1, 5, 7], nums2 = [2, 3, 4, 9],一共7个数,我们要找的是排在第4位的数。我们随便取一个数出来,比如nums1中的5,怎么判断它是不是排在第4呢?既然排第4,说明它前面有3个比它小的数。nums1里它排第2,所以只有1个比它小的数,那么显然nums2中必须有2个比它小的数才符合要求。好,那我们就拿nums2[1]和nums2[2]和5比一比,如果nums2[1] <= 5并且nums2[2] >= 5,5就是中位数,否则就不符合条件。不幸的是,nums2[1]和nums2[2]都比5小,说明我们取出来的这个数太大了,应该在5的左边找一个更小的数来试一试,这时候就可以用二分查找。后面采用递归的方式就可以一步一步逼近最终结果。

其实这个思想在网上有很多人提过,但是算法的实现大多是漏洞百出,经常有数组越界、数组中有重复元素时无法获得正确结果、移动过量导致程序崩溃等等,主要是各种各样的临界情况没有考虑周全。花了点时间整理了一个可以跑通leetcode测试的代码:

public class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        // handle boundry scenarios
        if (nums1.length == 0 && nums2.length == 0) {
            return 0;
        } else if (nums1.length == 0) {
            return calculateMedian(nums2);
        } else if (nums2.length == 0) {
            return calculateMedian(nums1);
        }
        
        return findMedian(nums1, nums2, 0, nums1.length - 1);
    }
    
    private double findMedian(int[] nums1, int[] nums2, int start, int end) {
        int middle = (nums1.length + nums2.length - 1) / 2; 
        int current = (start + end) / 2;
        
        // if median not in nums1, search in nums2
        if (start > end) {
            return findMedian(nums2, nums1, 0, nums2.length - 1);
        }
        
        // over moved, roll back
        // e.g. [1, 2, 3, 5, 6, 7] [4], middle is 3, current is (3+5)/2 = 4
        if (middle < current) {
            return findMedian(nums1, nums2, start, current - 1);
        }
        
        // not enough small numbers in nums2, need to increase current
        if (middle - current > nums2.length) {
            return findMedian(nums1, nums2, current + 1, end);
        }
        
        // hit condition: all numbers less than median are in nums1
        if (middle == current) {
            if (nums1[current] <= nums2[0]) {
                // Bingo! Found the median here
                return calculateMedian(nums1, nums2, current, 0);
            } else {
                // current value too large, move forward
                return findMedian(nums1, nums2, start, current - 1);
            }
        }
        
        // hit condition: all numbers in nums2 are less than median
        if (middle - current == nums2.length) {
            if (nums1[current] == nums2[nums2.length - 1]) {
                // Bingo! Found the median here
                return calculateMedian(nums1, nums2, current, nums2.length-1);
            } else if (nums1[current] > nums2[nums2.length - 1]) {
                // Bingo! Found the median here, -1 means next value locates in nums1
                return calculateMedian(nums1, nums2, current, -1);
            } else {
                // current value too small, move backward
                return findMedian(nums1, nums2, current + 1, end);
            }
        }
        
        // hit condition: has "middle - current" numbers in nums2 not greater than current
        if (nums1[current] >= nums2[middle-current-1] && nums1[current] <= nums2[middle-current]) {
            // Bingo! Found the median here
            return calculateMedian(nums1, nums2, current, middle-current);
        } else if (nums1[current] < nums2[middle-current-1]) {
            // current value too small, move backward
            return findMedian(nums1, nums2, current + 1, end);
        } else {
            // current value too large, move forward
            return findMedian(nums1, nums2, start, current - 1);
        }
    }
    
    private double calculateMedian(int[] nums) {
        int m = nums.length / 2;
        return (nums.length % 2 != 0 ? nums[m] : ((double)(nums[m-1] + nums[m])) / 2);
    }
    
    private double calculateMedian(int[] nums1, int[] nums2, int current, int nextInNums2) {
        int totalLen = nums1.length + nums2.length;
        if (totalLen % 2 != 0) {
            return nums1[current];
        } else if (current < nums1.length - 1) {
            int next = nextInNums2 < 0 ? nums1[current+1] : Math.min(nums1[current+1], nums2[nextInNums2]);
            return ((double)(nums1[current] + next)) / 2;
        } else {
            return ((double)(nums1[current] + nums2[nextInNums2])) / 2;
        }
    }
}

你可能感兴趣的:(leetcode)