[LeetCode]2 Median of Two Sorted Arrays (C++，Python实现)

LeetCode OJ的第二题，如果有问题或者给我指点欢迎来信讨论ms08.shiroh@gmail.com

题目描述

LeetCode OJ的第二题，题目描述如下：

There are two sorted arrays A and B of size m and n respectively. 

Find the median of the two sorted arrays. 

The overall run time complexity should be O(log (m+n)).

给定两个排序后的序列A和B， 大小分别为m， n，求这两个序列的中位数，也就是把这两个序列合并起来的大的序列的中位数。要求算法复杂度为O(log(m+n))

思路

问题的难点在于两个序列。如果一个已经排好序的序列，取中位数就是O(1)的操作。而且中位数为当变量值的项数N为奇数时，处于中间位置的变量值即为中位数；当N为偶数时，中位数则为处于中间位置的2个变量值的平均数。

而如果要把两个序列merge起来，那么时间复杂度达不到要求的O(lg(m+n))

看到O(lg(m+n))加上查找，我们能够很自然地想到二分查找。但是二分查找在两个序列就行不通了。但是仍然可以借鉴其思想。两个已经排好序的序列A, B，从中各取一个数，设分别取index1, index2个元素（编程时候要注意下标是从0开始，而这里的index是从1开始），假设A[index1]<B[index2]，那么能够得到一个结论，如下图

图中橙色部分肯定比绿色和红色部分要小，而图中红色部分肯定比蓝色和橙色部分要大。由于我们求的是整体的mid，而如果我们取index1=mid/2向下取整，index2=mid/2向下取整，那么橙色的部分加上蓝色的部分之和肯定小于mid，所以橙色部分必定在小于mid的集合中（因为即使橙色部分都比蓝色部分大，也在小于mid范围内，而蓝色部分可能比绿色部分大，所以不能够去除），可以删去橙色部分，并修改mid的值为mid-index1。如此往复，最后必然会有mid=1（因为到mid=1的时候mid>>1为0，不会再减少了），这时候可以根据是奇数还是偶数返回需要的值。

返回值是有技巧的，因为如果是偶数，那么要找处于中间位置的2个变量值mid, mid+1，而mid+1过程就包括了寻找mid的过程，所以可以一起做了而不用分开。

时间复杂度上，由于每次去除掉mid/2部分（在mid/2小于数组长度的情况下），所以可以分析出时间复杂度为O(lg(m+n))，符合题目要求。

代码

Python:

class Solution:
    # @return a float
    def findMedianSortedArrays(self, A, B):
        if len(A) + len(B) <= 1:
            return A[0] if not B else B[0]
        mid = (len(A) + len(B) + 1) >> 1
        (val1, val2) = self.search(A, 0, B, 0, mid)
        if (len(A) + len(B)) % 2 == 1:
            return val1
        else:
            return (val1 + val2) / 2.0

    # find mid and mid + 1
    # m means the begin index of A, n means the begin index of B
    def search(self, A, m, B, n, mid):
        if m >= len(A):
            return (B[n+mid-1], B[n+mid])
        if n >= len(B):
            return (A[m+mid-1], A[m+mid])
        if mid == 1:
            if A[m] < B[n] and m < len(A)-1:
                return (A[m], min(A[m+1], B[n]))
            if A[m] > B[n] and n < len(B)-1:
                return (B[n], min(A[m], B[n+1]))
            return (min(A[m], B[n]), max(A[m], B[n]))

        tmp_mid = mid >> 1
        index1 = min(m+tmp_mid, len(A))
        index2 = min(n+tmp_mid, len(B))
        if A[index1-1] < B[index2-1]:
            mid = mid - index1 + m
            m = index1
        else:
            mid = mid - index2 + n
            n = index2
        return self.search(A, m, B, n, mid)

C++

class Solution {
public:
    double findMedianSortedArrays(int A[], int m, int B[], int n) {
	// 因为可能在m+n为偶数时候一次处理两个数，所以要先排除掉只有一个元素的情况
        if (m + n == 1)
	    return m > 0 ? A[0] : B[0];
	return search(A, 0, m, B, 0, n);
    }

    // index_A表示当前A处理的开始下标，index_B同理。函数一次性处理长度和为偶数，需要算两个数的情况
    double search(int A[], int index_A, int m, int B[], int index_B, int n){
	int mid = ((m + n + 1) >> 1) - index_A - index_B;
	// A已经为空
	if (index_A >= m)
	    return (m+n) % 2 == 1 ? B[index_B+mid-1] : (B[index_B+mid-1] + B[index_B+mid]) / 2.0;
	    // B已经为空
	    if (index_B >= n)
		return (m+n) % 2 == 1 ? A[index_A+mid-1] : (A[index_A+mid-1] + A[index_A+mid]) / 2.0;
	    // mid降到了1，可以寻找解了
	    if (mid == 1){
		if (A[index_A] < B[index_B] && index_A < m-1)
		    return (m+n) % 2 == 1 ? A[index_A] : (A[index_A] + min(A[index_A+1], B[index_B])) / 2.0;
		if (B[index_B] < A[index_A] && index_B < n-1)
		    return (m+n) % 2 == 1 ? B[index_B] : (B[index_B] + min(B[index_B+1], A[index_A])) / 2.0;
		return (m+n) % 2 == 1 ? min(A[index_A], B[index_B]) : (A[index_A] + B[index_B]) / 2.0;
	    }
	    mid = mid >> 1;
	    int index1 = min(mid+index_A, m);
	    int index2 = min(mid+index_B, n);
	    if (A[index1-1] < B[index2-1])
		index_A = index1;
	    else
		index_B = index2;
	    return search(A, index_A, m, B, index_B, n);
	}

	int min(int a, int b){
	    return a > b ? b : a;
	}

	int max(int a, int b){
	    return a > b ? a : b;
	}
};