将“寻找中位数”转化成更一般化的“寻找第k大的数”这个问题。因为数组已经排序,所以可以使用二分缩小范围,使得时间复杂度满足要求。
假设A、B数组里都至少有k/2个元素,那么可以将k平均划分成k/2和k/2两半,在A里找第k/2大的元素,在B里找k/2大的元素(对k划分,解题的关键)。因为数组已排序,所以可以立即得到A中第k/2大的元素是A[k/2],B中第k/2大的元素是B[k/2]。然后将这两个数做比较。共有三种情况:
1. A[k/2] < B[k/2]
2. A[k/2] = B[k/2]
3. A[k/2] > B[k/2]
先来看第二种情况。如果A[k/2]=B[k/2],那么不用找了,A[k/2]或B[k/2]就是第k大的数。为什么?想象一下现在将A、B数组合并,那么A[k/2]前面最多可以放(k/2-1)+(k/2-1)+1=k-1个数(最后的1代表的是B[k/2]),所以A[k/2]就是第k大的数,同理对于B[k/2]也是一样。
再来看剩下的一、三两种情况,实际上他们是一样的。以第一种情况为例,此时A[k/2] < B[k/2],那么我们可以得出结论,第k大的数肯定不在A[0~k/2]里面。为什么?还是想象一下将A、B数组合并,A[k/2]前面最多可以放(k/2-1)+(k/2-1)=k-2个数,所以A[k/2]最多是第k-1大的数,那么在A[k/2]之前的数(A[0~k/2])更不可能是第k大的数了。所以在这种情况下,可以直接排除A中一半的元素,然后在剩余的数里继续寻找。
如果A或B非空但元素数量不够k/2,那么就不能平均划分了。不妨假设A中元素数量m < k/2,则划分成m和k-m两半,即在A中找第m大的数,在B中找第k-m大的数,同理由于数组已经排序,所以最后得到的数就是A[m]和B[k-m]。然后将这两数作比较,之后的处理过程跟上面一样。
如果A或B为空,那么更简单了,直接在另一个数组中返回第k大的数即可。
代码:
1 int findKth(int A[], int m, int B[], int n, int k) { 2 if (m > n) // 保证A数组元素个数小于B数组元素个数 3 return findKth(B, n, A, m, k); 4 if (m == 0) 5 return B[k - 1]; 6 if (k == 1) // 处理K等于1的特殊情况,否则第12行s-1=-1,数组越界 7 return min(A[0], B[0]); 8 9 int s = min(k / 2, m); 10 int t = k - s; 11 12 if (A[s - 1] == B[t - 1]) 13 return A[s - 1]; 14 else if (A[s - 1] < B[t - 1]) 15 return findKth(A + s, m - s, B, n, k - s); 16 else 17 return findKth(A, m, B + t, n - t, k - t); 18 } 19 20 double findMedianSortedArrays(int A[], int m, int B[], int n) { 21 if ((m + n) % 2 == 0) 22 return (findKth(A, m, B, n, (m + n) / 2) + findKth(A, m, B, n, (m + n) / 2 + 1)) / 2.0; 23 else 24 return findKth(A, m, B, n, (m + n) / 2 + 1); 25 }
关于代码再补充一个小技巧:数组索引是从0开始的,但是数组长度是从1开始的,再加上一些边界条件,很容易就弄混了。我的做法是,统一采用一种计数规范。比如,统一让所有索引变量从0开始计数,或统一让所有索引变量从1开始计数。
在上面的代码中,题目给的参数(20行)m和n表示数组长度,是从1开始计数的,所以之后的代码中出现的所有参数:m、n、k、s、t全都是从1开始计数的,当要访问数组元素的时候,统一减去1(如第5行、第12行等)。
时间复杂度
函数`findKth`每次执行的时间代价是O(1),当每次都对半分的时候,函数栈最深,执行次数最多,有log(m+n)次,所以时间复杂度是O(log(m+n))。
空间复杂度
函数`findKth`每次执行的空间代价是O(1),同理,函数栈深度最多是O(log(m+n)),所以空间复杂度是O(log(m+n))。