LeetCode OJ Median of Two Sorted Arrays

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)).

刚开始的时候看到discuss里有人说可以通过findKthSmallest问题来解决（在两个有序的数组中找到第k小的数），于是在

http://leetcode.com/2011/01/find-k-th-smallest-element-in-union-of.html看到了解析就用上了：

大概翻译一下内容：

维持一个等式：i + j = k - 1，其中i是A数组里的第i个数，j是B数组里的第j个数，k是所要求得的第k小的数。注意i和j都是下标（从0开始，小标为-1会特殊处理），而k是第k小，从1开始。

那么：

若找到Bj - 1 < Ai <= Bj，那么Ai即是第k小的数，因为Ai在A数组之前有i个数比Ai小，而Bj在B数组中有j个数比Bj小，那么对于两个数组来说，Ai之前总共就有i + j个数比Ai小，由 i + j = k - 1，Ai前有k - 1个数比其小，可知Ai即为第k小的数

若找到Ai - 1 < Bj <= Ai，那么Bj即是第k小的数，理由类似。

若上述两个关系没有成立，那么就继续以ij为点对AB数组进行二分，直到上述两个关系成立。划分的方法如下：

对于Ai和Bj：

若Ai <= Bj，那么必有Ai <= Bj - 1，否则就是 Bj - 1 < Ai <= Bj成立了。那么Ai在A数组里有i个比其小的数在前面，而在B数组中Ai最多只有j - 1个数比其小（Bj - 1前有j - 1个数），那么在两个数组中Ai前面就有i + j - 1个数比其小了，由 i + j = k - 1，Ai不可能是第k小的数，那么在A数组中比Ai小的数就更不可能了。因此，查找第k小的数的范围应该是排除了A数组中前i + 1个数（Ai前面有i个数，加上Ai本身一共i + 1个）。

若Bj < Ai，同理有Bj <= Ai - 1，分析与上类似，查找范围应该是排除了B数组中前j + 1个数。

class Solution {
public:
    double findMedianSortedArrays(int A[], int m, int B[], int n) {
        if ((m + n) & 1)  // (m + n) is odd
            return findKthSmallest(A, m, B, n, (m + n) / 2 + 1);
        else
            return (findKthSmallest(A, m, B, n, (m + n) / 2) + findKthSmallest(A, m, B, n, (m + n) / 2 + 1)) / 2.0;
    }
    int findKthSmallest(int A[], int m, int B[], int n, int k) {
        int i = (int)((double)m / (m + n) * (k - 1));
        int j = (k - 1) - i;

        // Note: A[-1] = -INF and A[m] = +INF to maintain invariant
        int Ai_1 = ((i == 0) ? INT_MIN : A[i-1]);  // when i == 0, A[i - 1] is INT_MIN
        int Bj_1 = ((j == 0) ? INT_MIN : B[j-1]);  // when j == 0, B[j - 1] is INT_MIN
        int Ai   = ((i == m) ? INT_MAX : A[i]);    // when i == m, A[i] is INT_MAX
        int Bj   = ((j == n) ? INT_MAX : B[j]);    // when j == n, B[j] is INT_MAX

        if (Bj_1 < Ai && Ai <= Bj) return Ai;  // found
        if (Ai_1 < Bj && Bj <= Ai) return Bj;  // found

        if (Ai <= Bj) return findKthSmallest(A + i + 1, m - (i + 1), B, j, k - (i + 1));
        else         return findKthSmallest(A, i, B + j + 1, n - (j + 1), k - (j + 1));
    }
};

后来自己试着用k的复杂度自己做了一次，很多情况要考虑：

class Solution {
public:
    double findMedianSortedArrays(int A[], int m, int B[], int n) {

        if (m == 0) {  // if A's size is zero
            if (n & 1) return B[n / 2];
            else return (B[n / 2 - 1] + B[n / 2]) / 2.0;
        }
        if (n == 0) {  // if B's size is zero
            if (m & 1) return A[m / 2];
            else return (A[m / 2 - 1] + A[m / 2]) / 2.0;
        }

        if ((m + n) & 1)  // (m + n) is odd
            return findKthSmallestForOdd(A, m, B, n, (m + n) / 2 + 1);
        
        // (m + n) is even
        int Apos, Bpos, k;
        k = (m + n) / 2;
        Apos = Bpos = -1;
        while (Apos + Bpos < k - 2) {
            if (Apos == m - 1) {
                Bpos++;
                continue;
            } else if (Bpos == n - 1) {
                Apos++;
                continue;
            }
            if (A[Apos + 1] <= B[Bpos + 1]) Apos++;
            else Bpos++;
        }

        int first;  // the first number in the pair of the median
        int firstPos;  // the position of the first number
        bool firstInA;  // if the first number is in A
        if ((Apos == -1) || (Bpos == -1)) {
            if (Bpos == -1) {  // if the k numbers are all in A
                first = A[Apos];
                firstPos = Apos;
                firstInA = true;
            }
            if (Apos == -1) {  // if the k numbers are all in B
                first = B[Bpos];
                firstPos = Bpos;
                firstInA = false;
            }
        } else {  // if the k numbers are in A and B
            first = A[Apos] >= B[Bpos] ? A[Apos] : B[Bpos];
            firstPos = A[Apos] >= B[Bpos] ? Apos : Bpos;
            firstInA = A[Apos] >= B[Bpos] ? true : false;
        }
        if (firstInA) {
            if (firstPos == m - 1) {  // if the first number is at the end of A, the second one must be at Bpos + 1
                return (first + B[Bpos + 1]) / 2.0;
            } else {
                if (Bpos == n - 1) return (first + A[firstPos + 1]) / 2.0;  // if the number in B are used up, the second number is behind the first one in A
                int second = B[Bpos + 1] < A[firstPos + 1] ? B[Bpos + 1] : A[firstPos + 1];  // else, the second number is the smaller one between the numbers behind the first one and the Bpos
                return (first + second) / 2.0;
            }
        } else {
            if (firstPos == n - 1) {
                return (first + A[Apos + 1]) / 2.0;
            } else {
                if (Apos == m - 1) return (first + B[firstPos + 1]) / 2.0;
                int second = A[Apos + 1] < B[firstPos + 1] ? A[Apos + 1] : B[firstPos + 1];
                return (first + second) / 2.0;
            }
        }
    }
    int findKthSmallestForOdd(int A[], int m, int B[], int n, int k) {
        int Apos, Bpos;
        Apos = Bpos = -1;
        while (Apos + Bpos < k - 2) {
            if (Apos == m - 1) {  // if the number in A are used up
                Bpos++;
                continue;
            } else if (Bpos == n - 1) {  // if the number in B are used up
                Apos++;
                continue;
            }
            if (A[Apos + 1] <= B[Bpos + 1]) Apos++;
            else Bpos++;
        }
        if (Bpos == -1) {  // if the k numbers are all in A
            return A[Apos];
        }
        if (Apos == -1) {  // if the k numbers are all in B
            return B[Bpos];
        }
        return A[Apos] > B[Bpos] ? A[Apos] : B[Bpos];
    }
};

当然还可以两种方法结合起来，也就是在第一种方法的(m + n)为偶数的情况中先用第一种方法找到组成中位数的前一个数，然后再用第二种方法的思想找到后一个数。比较复杂我还是不做了。