[LeetCode题解]从两个有序数组的并集中寻找第k小元素

Given two sorted arrays A, B of size m and n respectively. Find the k-th smallest element in the union of A and B. You can assume that there are no duplicate elements.

不得不承认这道题目解决起来非常的巧妙。像大多数难题一样，需要经过非常巧妙的观察才可以用简洁的方式求解。

朴素解法, O(m+n):

将两个数组进行合并，然后寻找第k小的元素可能非常直观。合并操作需要花费额外的O(m + n)的空间。线性运行时间已经很好了，但我们还能再做一些优化吗？

朴素解法的优化, O(k):

上面的方法可以这样优化，在此感谢提出该方法的读者Martin。

int findKthSMallest(int[] A, int[] B, int k) {int a_offset = 0, b_offset = 0;if (A.length + B.length < k) return -1;while (true) {if (a_offset < A.length) {while (b_offset == B.length ||A[a_offset] <= B[b_offset]) {a_offset++;if (a_offset + b_offset == k) return A[a_offset];}}if (b_offset < B.length) {while (a_offset == A.length ||A[a_offset] >= B[b_offset]) {b_offset++;}if (a_offset + b_offset == k) return B[b_offset];}}
}

使用两个指针，可以无需合并，遍历两个数组，因而不需要额外的空间。两个指针初始分别指向A和B的头，每一次将两个指针指向的数字中较小的那个的指针加一。第k小的数字总计遍历k次即可获得。该算法非常类似于寻找两个有序数组的交集。

最佳解法, O(log m + log n):

尽管上面的解法在时间和空间复杂度上都有了提升，但是仍然只能处理较小的k值，并且仍然是线性时间算法。我们还能再做优化吗？

上面的对数复杂度给了我们一个重要的提示。二分查找是对数复杂度算法的很好的例子，每次迭代时将查找空间折半。因此，为了达到O(logm + logn)的复杂度，我们必须在每轮迭代时将A和B的查找空间折半。

我们可以从比较A和B的中间元素出发解决这个比较棘手的问题，我们将这两个元素记为Ai和Bj。如果Ai在Bj和Bj-1之间，我们就恰好找到了第i+j+1小的元素。想想这是为什么。因此，如果我们选出i和j，使得i + j = k - 1，我们就可以找到第k小的元素。这是解决此问题的一个重要的等式变形。

总结一下上面的思路，

维护等式，
i + j = k – 1,
如果 Bj-1 < Ai < Bj, 那么 Ai 就是第k小的元素,
否则，如果 Ai-1 < Bj < Ai, 那么 Bj 就是第k小的元素

如果满足了上面的条件之一，我们就完成了问题的求解。如果没有满足，我们再使用i和j作为枢轴对数组进行划分。但是怎样做呢？我们应该舍弃哪一部分？还有Ai和Bj自己应当怎样处理？

我们观察一下，如果Ai < Bj，则 Ai < Bj-1一定成立。另一方面，如果Bj < Ai，则 Bj < Ai - 1。想想为什么呢？（因为迭代还没有结束，否则就直接返回Ai或者Bj了）

使用上面的关系式，可以明确当Ai < Bj时， Ai和它左边的元素肯定不会是第k小的元素。Bj及其右侧的元素也是如此。因此，我们可以直接舍弃Ai及其左边的元素，Bj及其右边的元素。

如果你还没有想通为什么这是正确的，试着画两个方块代表A和B。然后尝试可视化的将Ai左侧的块放在Bj-1之前。你应该可以很容易的看到插入块中不可能有元素是第k小的。对于后者，你可以结合不变式i + j = k - 1思考为什么Bj及其右侧的元素也不会是第k小的元素。

另一方面，Ai > Bj的情况可以依次类推。很容易。

下面的代码中我加入了很多断言（非常推荐这种编程方式）来帮助你理解代码。注意下面的代码采用了尾递归，因此你可以很容易地改写成迭代的形式。不过，我在这里不做详细描述，因为这就是我解决问题的思路，通过递归的方式表达可能更加自然。

另外的旁注考虑了i和j的选择。下面的代码将数组的大小作为权重对数组进行划分。这样做的原因是可以更快的找到所求元素（除非极端情况下A的所有元素都小于B）。其实你可以自行选择i作为A的中点。理论上，你可以选择任何满足i + j = k - 1条件的i值和j值。

int findKthSmallest(int A[], int m, int B[], int n, int k) {assert(m >= 0); assert(n >= 0); assert(k > 0); assert(k <= m+n);int i = (int)((double)m / (m+n) * (k-1));int j = (k-1) - i;assert(i >= 0); assert(j >= 0); assert(i <= m); assert(j <= n);// invariant: i + j = k-1// Note: A[-1] = -INF and A[m] = +INF to maintain invariantint Ai_1 = ((i == 0) ? INT_MIN : A[i-1]);int Bj_1 = ((j == 0) ? INT_MIN : B[j-1]);int Ai   = ((i == m) ? INT_MAX : A[i]);int Bj   = ((j == n) ? INT_MAX : B[j]);if (Bj_1 < Ai && Ai < Bj)return Ai;else if (Ai_1 < Bj && Bj < Ai)return Bj;assert((Ai > Bj && Ai_1 > Bj) || (Ai < Bj && Ai < Bj_1));// if none of the cases above, then it is either:if (Ai < Bj)// exclude Ai and below portion// exclude Bj and above portionreturn findKthSmallest(A+i+1, m-i-1, B, j, k-i-1);else /* Bj < Ai */// exclude Ai and above portion// exclude Bj and below portionreturn findKthSmallest(A, i, B+j+1, n-j-1, k-j-1);

原文链接：http://leetcode.com/2011/01/find-k-th-smallest-element-in-union-of.html