【Leetcode】4. Median of Two Sorted Arrays

题目地址：

https://leetcode.com/problems/median-of-two-sorted-arrays/

给定两个升序数组$A$和$B$，求两个数组所有数的中位数。如果一共偶数个，则取最中间的两个数的平均值。

设各自的长度$n$和$m$，那么如果我们解决了对任意给定的$k$找第$k$小的数，这个问题就解决了。所以先考虑如何解决这个更加一般的问题。我们考虑数组$A$和$B$的各自的前$k/2$（以下的除法都是整除）个数。如果$A[k/2-1]<B[k/2-1]$，那么我们可以断言$A[0],...,A[k/2-1]$一定都不是答案，因为如果存在$j\le k/2-1$使得$A[j]$是第$k$小的数，由于$A[k/2-1]<B[k/2-1]$，所以比$A[j]$还小的数的个数是小于等于$j+k/2-1<k$的，这就矛盾了。所以我们就可以将$A[0],...,A[k/2-1]$舍去，直接从$A[k/2]$开始，加上$B$的所有数，找第$k-k/2$小的数。这样就将要找第$k$小的数转变为了要找第$k-k/2\le k/2+1$小的数，而且重复这种操作时，$k$必然会严格下降到$1$，而对于$k=1$答案是显然的。如果$A$的长度是小于$k/2$的，那么就可以直接断定$B[0],...,B[k/2-1]$不可能是答案，可以直接从$B[k/2]$开始和$A$全体中寻找答案，剩下的逻辑与上面类似。代码如下：

class Solution {
 public:
  double findMedianSortedArrays(vector<int>& v1, vector<int>& v2) {
    int n = v1.size() + v2.size();
    if (n % 2)
      return find(v1, 0, v2, 0, n / 2 + 1);
    else
      return (find(v1, 0, v2, 0, n / 2) + find(v1, 0, v2, 0, n / 2 + 1)) / 2.0;
  }

  int find(vector<int>& v1, int i, vector<int>& v2, int j, int k) {
    if (v1.size() - i > v2.size() - j) return find(v2, j, v1, i, k);
    if (v1.size() == i) return v2[j + k - 1];
    if (k == 1) return v1.size() == i ? v2[j] : min(v1[i], v2[j]);
    int k1 = min((int)v1.size() - 1, i + k / 2 - 1),
        k2 = j + k - (k1 - i + 1) - 1;
    if (v1[k1] < v2[k2])
      return find(v1, k1 + 1, v2, j, k - (k1 - i + 1));
    else if (v1[k1] > v2[k2])
      return find(v1, i, v2, k2 + 1, k - (k2 - j + 1));
    else
      return v1[k1];
  }
};

时间复杂度$O(\log (m+n))$，空间$O(\log k)$（递归栈深度）。

时间复杂度证明：
设$A$和$B$的数字个数和为$N$，则$$\begin{gathered} T(N,k)=T(N-k/2,k/2)+O(1) \\ =T(N-k/2-k/4,k/4)+O(2) \\ =...=T(N-k,1)+O(\log k)=O(\log k)=O(\log N) \end{gathered}$$