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)).
这是一道很容易让人产生错觉的题。初看之下思路很清晰, 只要两个pointer连续比较后移就能得到O(n)的解法,而O(log(m+n))的解法只要通过二分法就行。以前写两个等长数列的题目并不吃力,误以为只要处理一些边界条件就能得到这题的答案。然而,仔细开始写之后发现不对劲:
- 等长的情况下,每次都可以从两个数列各自舍掉一半,很快就能得到结果。对于不等长的情况,各自舍掉一半的结果是新的中位数和原来的不一样。
- 等长情况下,元素总量为偶数,可以确保最后的中位数是两个数字的平均值。不等长后,两个数列的长度奇偶情况总共有四种组合。奇数个元素的中位数和偶数个元素的中位数比较后得到的结果分析起来更复杂。
- 假设每次从两个数组各自删除一部分,最终两个数列都可能先为空,边界条件判断比较繁复。
这些思路上的障碍和边界条件的判断组合起来后,整个程序思路中出现很多细节,代码冗长。
参考这里之后,才理清了思路,几个关键点:
- 将问题从中位数generalize到求两个数组中的第k个数。这样一来,最终要的中位数只取决于最开始的两个数组长度m, n。而对于中间过程而言,m与n各自的奇偶性不再重要,关注的重点从中位数这一具有不确定性(可能是一个数或者两个数的平均值)的目标变成了一个确定的数值:第k个值。
- 对A和B的长度做一个基本假设:size(A) <= size(B), 如果 A比B长,直接将A与B进行对换。这样一来,对于递归过程中m与n的边界条件判断变得更简洁
- 比较A[k/2 -1]与B[k + 1 – k / 2] 的值,如果二者相等,说明这一值就是中位数。否则,小值所在的数列中从头到小值的部分都可以确定小于第k个值,可以直接移除后递归寻找
- 在比较A[k/2 -1]与B[k + 1 – k / 2] 的时候,还需要考虑k与m,n的关系,即:k/2 – 1 >= m, 其中一个数列太短,无法取到想要的值,这时可以取最后一个值A[m-1]进行比较,只需要确认两部分比较的长度之和为k。由于这一情况最多在递归最内层出现一次,总的时间复杂度保持在O(log(m+n))
代码如下:
class Solution { public: double findKth(int A[], int m, int B[], int n, int k) { if (m > n) return findKth(B, n, A, m, k); if (m == 0) return B[k - 1]; if (k == 1) return std::min(A[0], B[0]); int la = std::min(k / 2, m); int lb = k - la; if (A[la - 1] == B[lb - 1]) return A[la - 1]; else if (A[la - 1] > B[lb - 1]) return findKth(A, m, B + lb, n - lb, k - lb); else return findKth(A + la, m - la, B, n, k - la); } double findMedianSortedArrays(int A[], int m, int B[], int n) { int s = m + n; if (s & 1) return findKth(A, m, B, n, s / 2 + 1); else return (findKth(A, m, B, n, s / 2) + findKth(A, m, B, n, s / 2 + 1)) / 2; } };