LeetCode #4. Median of Two Sorted Arrays プログラミング練習

問題:

There are two sorted arrays nums1 and nums2 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)).

Example 1:
nums1 = [1, 3]
nums2 = [2]

The median is 2.0
Example 2:
nums1 = [1, 2]
nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

難易度: hard
入力: 二つのint型Array
目的:  入力された数値の中央値を探せ
出力: int
注意: 入力された二つのArrayはソートされてあり、最適解はO(log(m+n))である

初hard問題です。
流石にhardだけあってボーダーケースがとても多く、普通に回答したらすぐにTLEになってしまいます。
注意すべき事: 

  • nums1の数字が全部nums2より小さい場合がある
  • nums1とnums2の長さの和が偶数の場合中央値の隣の数字も考慮しなければならない
  • 中央値の隣の数値が同じArrayの中にあるとは限らない
  • 二つのArrayを統合させソートし中央値を出すのが一番簡単だがTLEになる

問題文から最適解はO(log(m+n))だと言われているので、明らかに二分探索を使えと言うヒントが与えられている。
私がまず頭に浮かんだ方法はO(m+n)の回答ですがAcceptされたので貼り付けます。

public class Solution {  
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {  
        int l1 = nums1.length;  
        int l2 = nums2.length;  
        int mid = (l1+l2)/2;  
        int[] s = new int[mid+1];  
        int j=0,k=0;  
        for(int i=0;i<mid+1;i++){  
            if(j<l1&&(k>=l2||nums1[j]<nums2[k])){  
                s[i]=nums1[j++];  
            }else{  
                s[i]=nums2[k++];  
            }  
        }  
        if ((l1+l2)%2!=0){  
            return (double)s[mid];  
        }else{  
            return (double)(s[mid-1]+s[mid])/2;  
        }  
    }  
}  

二つのArrayを下から中央値まで統合させそのまま出力する方法ですね。
注意したのは偶数と奇数を分けるくらいかな?
問題文はO(log(m+n))で済ませと言っているのですが何故か通ってしまいました。
ちなみにrun timeは80 msで、上位26.6%でした
あれ割と高い(笑)

二分探索での良い案が思いつかなかったのでネットの答えを参考に書いてみました。
run time は 62 msで、上位11.2%でした。
以下、回答です。

public class Solution {  
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {  
        int m = nums1.length, n = nums2.length;  
        int k = (m + n) / 2;  
        if((m+n)%2==0){  
            return (findKth(nums1,nums2,0,0,m,n,k)+findKth(nums1,nums2,0,0,m,n,k+1))/2;  
        }   else {  
            return findKth(nums1,nums2,0,0,m,n,k+1);  
        }  
  
    }  
  
    private double findKth(int[] arr1, int[] arr2, int start1, int start2, int len1, int len2, int k){  
        if(len1>len2){  
            return findKth(arr2,arr1,start2,start1,len2,len1,k);  
        }  
        if(len1==0){  
            return arr2[start2 + k - 1];  
        }  
        if(k==1){  
            return Math.min(arr1[start1],arr2[start2]);  
        }  
        int p1 = Math.min(k/2,len1) ;  
        int p2 = k - p1;  
        if(arr1[start1 + p1-1]<arr2[start2 + p2-1]){  
            return findKth(arr1,arr2,start1 + p1,start2,len1-p1,len2,k-p1);  
        } else if(arr1[start1 + p1-1]>arr2[start2 + p2-1]){  
            return findKth(arr1,arr2,start1,start2 + p2,len1,len2-p2,k-p2);  
        } else {  
            return arr1[start1 + p1-1];  
        }  
    }  
}

中央値を探すのではなく、入力の中でk番目に大きい数値を探すに変換した考え方ですね。
nums1のp番目の数字とnums2のq番目の数字がk番目の数字(目的数値)より小さいとし、p+q=k-1の時、p+1又はq+1がk番目の数字になります。
短めのArrayをarr1とし、長めのArrayをarr2とします。
arr1のp番目の数値とarr2のq番目の数値を比較し、arr1の方が小さい場合、pより前の数値はk番目の数値より小さいので破棄します。
pから始まる新しいarr1とarr2の中でk-p番目に大きい数値を探します。
これを繰り返し、arr1が空になる場合、arr2からk番目の数値をreturnします。
k=1の場合、arr1かarr2の一番小さい数字がk番目の数値なので小さい方をreturnします。

結構複雑なアルゴリズムになってしまったのでもしわからないことがあればメッセージ下さい。
ですがこのレベルの問題が面接で出された場合、
経験上最初の答えでも割とパスされる事が多いですね(笑)

では