[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 invariant
  int 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 portion
    return findKthSmallest(A+i+1, m-i-1, B, j, k-i-1);
  else /* Bj < Ai */
    // exclude Ai and above portion
    // exclude Bj and below portion
    return 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

本文链接:http://bookshadow.com/weblog/2015/02/05/find-kth-smallest-element-in-the-union-of-two-sorted-arrays/
请尊重作者的劳动成果,转载请注明出处!书影博客保留对文章的所有权利。

如果您喜欢这篇博文,欢迎您捐赠书影博客: ,查看支付宝二维码