题目: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)).
刚开始的时候理解有误,以为是求中位数。其实不是。这里所谓的"median of the two sorted arrays"指的是:如果A.length + B.length 是奇数,返回merge后中间的那个数;如果是偶数,返回中间两个数的平均数。
例如:[1,2]和[3],返回2;而[1,2]和[3,4],返回2.5
解法:不难有O(m+n)的做法,即merge A 和 B,得到一个长的sorted数组,返回中间一个或中间两个的平均数。但是不符合题目要求。
如何做到O(log(m+n))呢?
首先考虑如何找到在A和B数组merge的第k个元素。可以对k进行分治。假设A的长度为m,B的长度为n,m<n,A的offset为aoffset,B的offset为boffset,初始为0
- 如果m为0,返回B[k-1];
- 如果k是1,返回A[aoffset]和B[boffset]中比较小的那个;
- 我们从A中取前(相对于其offset的)Ka个(一般是k/2),从B中取前Kb个(Ka + Kb = k);
- 比较A[aoffset + Ka-1]和B[boffset + Kb-1]。如果A[aoffset + Ka-1] 大于B[boffset + Kb-1],说明要找的第k个元素肯定不在B中刚刚选取的Kb个元素中,否则的话,在选择到的总共K个数中,最大的数应该在B中(但是实际在A中);反之,第k个元素肯定也不会在A中选取的Ka个元素中。也可以这样想,通过比较A[aoffset + Ka-1]和B[boffset + Kb-1],我们知道在第3步中对K值的划分是否足够“偏向”目标元素,从而排除掉一部分元素;
- 如果A[aoffset + Ka-1] 大于B[boffset + Kb-1],移动boffset到boffset+kb,同时k减去Kb。回到1.(相当于继续在A和新的B中寻找第k-Kb个元素)
- 如果A[aoffset + Ka-1] 小于等于B[boffset + Kb-1],移动aoffset到aoffset+ka,同时k减去Ka。回到1.
代码如下:
1 private static int findKth(int A[], int aoffset, int m, int B[],int boffset, int n,int k){ 2 if(m > n) return findKth(B, boffset, n, A, aoffset, m, k); 3 if(m==0)return B[k-1]; 4 if(k==1)return Math.min(A[aoffset], B[boffset]); 5 int pa = Math.min(k/2, m); 6 int pb = k - pa; 7 if(A[aoffset + pa - 1] >= B[boffset + pb - 1]) 8 return findKth(A,aoffset, m, B, boffset + pb, n - pb, k-pb); 9 else return findKth(A, aoffset + pa, m - pa, B, boffset, n, k - pa); 10 }
注意在递归调用之前,始终保持m<n。
有了这个方法,我们可以用它来计算两个有序数组的第K个元素,当然也包括中间的一个(或者两个的平均值):
1 public static double findMedianSortedArrays(int A[], int B[]) { 2 // Start typing your Java solution below 3 // DO NOT write main() function 4 //return findMiddleValue(A,B);//this is o(m+n) 5 int n = A.length + B.length; 6 if(n%2 == 0){ 7 return (double)(findKth(A,0,A.length, B,0,B.length,n/2) +findKth(A,0,A.length, B,0,B.length,n/2 + 1))/(double)2; 8 }else { 9 return findKth(A,0,A.length, B,0,B.length,n/2 + 1); 10 } 11 }