题目
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)).
解答
求中位数需要根据数组长度是奇数还是偶数分别讨论,奇数长度时中位数为最中间的一个数,偶数长度时中位数为最中间的两个数的平均值,为了方便,可以实现一个比题目更一般化的函数,求A和B的第k小数的函数,那么中位数的问题很容易解决。
求一个有序数组的第k个数只需要O(1)的复杂度,现在有两个数组,显然花费额外空间以O(n)时间归并然后O(1)寻找不满足题目要求。
既然要求log时间复杂度,一般需要使用到二分思想。分别考虑A和B的第k/2个元素:
- 如果它们相等,则第k个数为其中的任意一个
- 如果A中的比较大,则B中前k/2个元素都不可能是第k个数了,因为这个数至少应该为A的第k/2个数,把B的前k/2去掉,然后重新寻找。
- 如果B中的比较大,则把A的前k/2个数去掉,重新寻找。
直到A和B中某个变为空时或者寻找第1个数时可以停止递归,直接找到结果。
注意,上面的k/2只是理想的简单情况,实际上A和B的长度可能不够k/2,或者k为奇数等,但这些不是主要问题,可以让A取第k/2个数字,然后A不够长,则取A的最后一个数字,然后B取剩下长度对应的那个数字,具体参考代码。
代码如下: