LeetCode 4. Median of Two Sorted Arrays

 

문제

정렬된 두 정수 배열 $nums1$, $nums2$ 가 주어진다.
두 배열의 중앙값을 $\log$ 시간복잡도로 구해보자.

풀이

문제를 살짝 변형해서, 두 배열의 $k$번째 원소를 구해보자.

$k$번째 원소를 어떻게 구할까?

아직은 모르지만, $k$번째 원소의 값을 $v$라고 하자.
어떤 값 $x$가 $v$보다 작다면 두 배열에서 인덱스가 $k$ 미만이다.
따라서 $x$를 매개 변수로 두고 이분 탐색을 진행하며 가능한 $x$의 범위를 빠르게 좁혀나갈 수 있다.

두 배열에서 어떤 값 $x$의 인덱스는 이분 탐색으로 빠르게 구할 수 있다.

코드

class Solution {
    public int upperBound(int[] arr, int x) {
        int s = 0, e = arr.length, m = (s+e)>>1;
        while(s < e) {
            if(arr[m] <= x) s = m+1;
            else e = m;
            m = (s+e)>>1;
        }
        return m;
    }

    public int indexOf(int[] nums1, int[] nums2, int x) {
        return upperBound(nums1, x) + upperBound(nums2, x) - 1;
    }

    public int findKthValue(int[] nums1, int[] nums2, int k) {
        int s = -1000000, e = 1000000, m = (s+e)>>1;
        while(s < e) {
            if(indexOf(nums1, nums2, m) < k) s = m+1;
            else e = m;
            m = (s+e)>>1;
        }
        return m;
    }

    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int N = nums1.length, M = nums2.length;

        double answer = findKthValue(nums1, nums2, (N+M-1)/2);
        if((N+M) % 2 == 0) answer = (answer + findKthValue(nums1, nums2, (N+M)/2)) / 2;

        return answer;
    }
}