[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)).

You may assume nums1 and nums2 cannot be both empty.

 

풀이

전형적인 binary search 문제이다.

 

$l1=nums1.size(), l2=nums2.size()$라고 하자.

$nums1$에서 binary search를 돌리는 데 찾는 것은 다음과 같다.

'$nums1[i]$보다 작거나 같은 숫자가 $nums1, nums2$ 합쳐서 $(l1+l2+1)/2$ 개 이상이 되는 최소한의 index i'

 

현재 $nums1$에서의 index를 $m_1$이라고 하면 해당하는 $nums1[m_1]$보다 작은 숫자가 $nums1$에는 $m_1$개 있을 것이다.

따라서 $nums2$에 $m_2=(m1+m2+1)/2 - m_1$개가 있어야 한다.

 

리스트의 상황을 간단히 그리면 다음과 같다.

nums1[0]    nums1[1]   ...   nums1[$m_1-1$] |||  nums1[$m_1$]  ... 

nums2[0]    nums2[1]   ...   nums2[$m_2-1$] |||  nums2[$m_2$]   ...

 

각 list가 정렬되어 있음은 주어진 조건이므로 우리는

$nums2[m_2-1] <= nums1[m_1]$과 $nums1[m_1-1] <= nums2[m_2]$만 체크하면 충분하다.

첫번째 조건을 만족하지 못하면 $nums1[m_1]$이 너무 작다는 뜻이니, $m_1$을 증가시켜주면 되고

두번째 조건을 만족하지 못하면 반대로 $nums1[m_1]$이 너무 크다는 뜻이니 $m_1$을 감소시켜주면 된다.

 

이 조건을 이용하여 binary search를 진행하여 중간값을 찾는다.

그렇게 하면 |||로 나누어진 구간 왼쪽과 오른쪽에 숫자가 반반씩 있게 되고,

홀짝을 비교하여 ||| 앞뒤에 있는 4개의 숫자를 단순계산하여 중간값을 찾을 수 있다.

 

Time Complexity : $O(log(n+m))$

 

코드

class Solution {
public:
    double findMedianSortedArrays(vector& nums1, vector& nums2) {
        vector& n1=(nums1.size()<nums2.size()?nums1:nums2);
        vector& n2=(nums1.size()<nums2.size()?nums2:nums1);
        
        int l1 = n1.size(), l2=n2.size();
        int l=0, r=l1, midnum = (l1+l2+1)>>1;
        while(l<=r)
        {
            int m1 = (l+r)>>1;
            int m2 = midnum-m1;
            if(m1>r && n2[m2-1]>n1[m1]) l = m1+1;
            else if( m1>l && n1[m1-1]>n2[m2]) r=m1-1;
            else
            {
                int ans;
                if(m1==0) ans=n2[m2-1];
                else if(m2==0) ans=n1[m1-1];
                else ans = max(n1[m1-1], n2[m2-1]);
                if((l1+l2)&1) return ans;
                
                int ans2;
                if(m1==l1) ans2 = n2[m2];
                else if(m2==l2) ans2= n1[m1];
                else ans2 = min(n1[m1], n2[m2]);
                return (ans+ans2)/2.0;
            }
        }
        return 0.0;
    }
};