[LeetCode] 004. Median of Two Sorted Arrays (Hard) (C++/Java/Python)

索引：[LeetCode] Leetcode 题解索引 (C++/Java/Python/Sql)

Github: https://github.com/illuz/leetcode

004.Median_of_Two_Sorted_Arrays (Hard)

链接：

题目：https://oj.leetcode.com/problems/Median-of-Two-Sorted-Arrays/

代码(github)：https://github.com/illuz/leetcode

题意：

求两个有序数组的中位数。要求复杂度是 O(log(n + m))。

分析：

两种思路：

1. 直接 merge 两个数组，然后求中位数，能过，不过复杂度是 O(n + m)。

2. 用二分的思路去做，这不好想，还要考虑到奇偶。可以转化思维，去求两个有序数组中的第 K 大数，这样就比较好想了。

代码：

1. Merge

C++:

class Solution {
public:
	double findMedianSortedArrays(int A[], int m, int B[], int n) {
		vector<int> C;
		int pa = 0, pb = 0;	// point of A & B

		while (pa < m || pb < n) {
			if (pa == m) {
				C.push_back(B[pb++]);
				continue;
			}
			if (pb == n) {
				C.push_back(A[pa++]);
				continue;
			}
			if (A[pa] > B[pb])
				C.push_back(B[pb++]);
			else
				C.push_back(A[pa++]);
		}
		if ((n + m)&1)
			return C[(n+m)/2];
		else
			return (C[(n+m)/2 - 1] + C[(n+m)/2]) / 2.0;
	}
};

2. 二分

C++:

class Solution {
	private:
		double findKthSortedArrays(int A[], int m, int B[], int n, int k) {
			if (m < n) {
				swap(n, m);
				swap(A, B);
			}
			if (n == 0)
				return A[k - 1];
			if (k == 1)
				return min(A[0], B[0]);

			int pb = min(k / 2, n), pa = k - pb;
			if (A[pa - 1] > B[pb - 1])
				return findKthSortedArrays(A, m, B + pb, n - pb, k - pb);
			else if (A[pa - 1] < B[pb - 1])
				return findKthSortedArrays(A + pa, m - pa, B, n, k - pa);
			else
				return A[pa - 1];
		}

	public:
		double findMedianSortedArrays(int A[], int m, int B[], int n) {
			if ((n + m)&1)
				return findKthSortedArrays(A, m, B, n, (n + m) / 2 + 1);
			else
				return (findKthSortedArrays(A, m, B, n, (n + m) / 2 + 1) +
						findKthSortedArrays(A, m, B, n, (n + m) / 2)) / 2.0;
		}
};

Java:

public class Solution {

    private double findKthSortedArrays(int A[], int astart, int aend,
                                       int B[], int bstart, int bend, int k) {
        int m = aend - astart, n = bend - bstart;
        if (m < n) {
            return findKthSortedArrays(B, bstart, bend, A, astart, aend, k);
        }
        if (n == 0)
            return A[astart + k - 1];
        if (k == 1)
            return Math.min(A[astart], B[bstart]);

        int pb = Math.min(k / 2, n), pa = k - pb;
        if (A[astart + pa - 1] > B[bstart + pb - 1])
            return findKthSortedArrays(A, astart, aend, B, bstart + pb, bend, k - pb);
        else if (A[astart + pa - 1] < B[bstart + pb - 1])
            return findKthSortedArrays(A, astart + pa, aend, B, bstart, bend, k - pa);
        else
            return A[astart + pa - 1];
    }

    public double findMedianSortedArrays(int A[], int B[]) {
        int m = A.length, n = B.length;
        if ((n + m) % 2 == 1)
            return findKthSortedArrays(A, 0, m, B, 0, n, (n + m) / 2 + 1);
        else
            return (findKthSortedArrays(A, 0, m, B, 0, n, (n + m) / 2 + 1) +
                    findKthSortedArrays(A, 0, m, B, 0, n, (n + m) / 2)) / 2.0;
    }
}

Python:

class Solution:
    def findKthSortedArrays(self, A, B, k):
        if len(A) < len(B):
            tmp = A
            A = B
            B = tmp
        if len(B) == 0:
            return A[k - 1]
        if k == 1:
            return min(A[0], B[0])

        pb = min(k / 2, len(B))
        pa = k - pb
        if A[pa - 1] > B[pb - 1]:
            return self.findKthSortedArrays(A, B[pb:], k - pb)
        elif A[pa - 1] < B[pb - 1]:
            return self.findKthSortedArrays(A[pa:], B, k - pa)
        else:
            return A[pa - 1]

    # @return a float
    def findMedianSortedArrays(self, A, B):
        if (len(A) + len(B)) % 2 == 1:
            return self.findKthSortedArrays(A, B, (len(A) + len(B)) / 2 + 1)
        else:
            return (self.findKthSortedArrays(A, B, (len(A) + len(B)) / 2) +
                self.findKthSortedArrays(A, B, (len(A) + len(B)) / 2 + 1)) / 2.0