[LeetCode] Median of Two Sorted Arrays

문제 출처

leetcode.com/problems/median-of-two-sorted-arrays/

Median of Two Sorted Arrays - LeetCode

문제

Given two sorted arrays nums1 and nums2 of size m and n respectively, return the median of the two sorted arrays.
Follow up: The overall run time complexity should be O(log (m+n)).

 

각 요소들이 정렬된 2개의 배열 nums1과 nums2가 주어진다.

두 배열의 중간값을 구해야 한다.

그리고 시간 복잡도가 O(log(m+n)) 이어야 한다는 추가 조건이 있다.

예시 1

Input: nums1 = [1,3], nums2 = [2]
Output: 2.00000
Explanation: merged array = [1,2,3] and median is 2.

 

두 배열의 개수 합이 홀수면 그냥 가운데 요소를 찾으면 된다.

예시 2

Input: nums1 = [1,2], nums2 = [3,4]
Output: 2.50000
Explanation: merged array = [1,2,3,4] and median is (2 + 3) / 2 = 2.5.

 

두 배열의 개수 합이 짝수면 가운데 두 요소의 중간값을 구하면 된다.

예시 3

Input: nums1 = [0,0], nums2 = [0,0]
Output: 0.00000

예시 4

Input: nums1 = [], nums2 = [1]
Output: 1.00000

 

nums1 배열이 빈 배열일 수 있다는 것을 보여준다.

예시 5

Input: nums1 = [2], nums2 = []
Output: 2.00000

 

nums2 배열이 빈 배열일 수 있다는 것을 보여준다.

제약 사항

nums1.length == m
nums2.length == n
0 <= m <= 1000
0 <= n <= 1000
1 <= m + n <= 2000
10^6 <= nums1[i], nums2[i] <= 10^6

 

각 배열의 개수는 0 ~ 1000개이고 요소가 가질 수 있는 값은 int 값 범위 내이다.

그리고 두 배열 모두 비어있는 상황은 없다.

풀이

먼저 문제를 어떻게 풀어야 할지 고민해 보자.

 

문제를 풀 때 가장 단순한(Brute Force) 방법부터 시작하는 것이 좋다.

 

그냥 두 배열을 다시 정렬한 후에 중간값을 구하면 될 것이다.

 

일반적으로 빠른 정렬을 사용하더라도 시간 복잡도는 O((m+n)log(m+n))이 될 것이다.

 

위의 문제에서 정렬을 좀 더 빠르게 한다면 O(m+n)이 나올 것이다.

 

병합 정렬(merge sort)에서 divide & conquer 후 combine하는 과정을 똑같이 따라 하면 된다.

 

두 배열의 요소들을 앞에서부터 하나씩 비교해 가며 더 작은 요소를 새로운 배열에 추가하여 정렬하는 방법 말이다.

 

그래도 문제에서 요구하는 시간 복잡도를 충족시키지 못한다.

 

시간 복잡도 조건으로 인해 문제에서 제시하는 풀이 방법은 정렬하지 말라는 것으로 해석할 수 있다.

 

어떻게 해야 할까?

 

문제의 조건으로부터 힌트를 얻을 때가 꽤 자주있다.

 

O(log(m+n))을 어디서 많이 보지 않았는가?

 

이진 탐색(binary search)의 시간 복잡도이다.

 

일단 여기까지만 생각하고 다시 문제로 돌아가 보자.

 

nums1 배열과 nums2 배열은 정렬되어 있다.

 

nums1 배열의 어느 지점(i)과 nums2 배열 어느 지점(j)을 기준으로 각각 둘로 나눈다면 아래와 같을 것이다.

 

$$ nums1[0], nums1[1], ... nums1[i - 1]　|　nums1[i], nums1[i + 1], ... nums1[m - 1] $$

$$ nums2[0], nums2[1], ... nums2[j - 1]　|　nums2[j], nums2[j + 1], ... nums2[n - 1] $$

 

이 때 nums1[0], nums1[1], ... nums1[i - 1]과 nums2[0], nums2[1], ... nums2[j - 1]을 왼쪽 부분이라고 하고

 

nums1[i], nums1[i + 1], ... nums1[m - 1]과 nums2[j], nums2[j + 1], ... nums2[n - 1]을 오른쪽 부분이라고 하자.

 

여기서 왼쪽 부분의 개수가 오른쪽 부분의 개수와 같고 (또는 1개 차이가 나고)

(조건 1)

 

왼쪽 부분에 있는 모든 요소가 항상 오른쪽 부분에 있는 모든 요소보다 값이 작으면

(조건 2)

 

두 배열을 합쳐 다시 정렬했을 때 중간에 위치할 요소가 nums1[i - 1], nums1[i], nums2[j - 1], nums2[j] 중에 있게 된다.

 

그래서 위 두 조건을 만족할 i와 j를 찾으면 되는 것이다.

 

첫 번째 조건을 먼저 식으로 표현하면 아래와 같다.

 

$$ [식1]　i + j = (m - i) + (n - j) $$

$$ [식2]　i + j = (m - i) + (n - j) + 1 $$

$$ [식3]　i + j + 1 = (m - i) + (n - j) $$

 

먼저 [식1]부터 보자.

 

왼쪽 항(i + j)은 왼쪽 부분의 개수를 나타내고 오른쪽 항((m - i) + (n - j))은 오른쪽 부분의 개수를 나타낸다.

 

왼쪽 항에서 i는 nums1 배열의 왼쪽 부분의 개수이고 j는 nums2 배열의 왼쪽 부분의 개수이다.

(헷갈리면 m과n, i와 j에 직접 값을 넣고 세어서 확인해 보자.)

 

오른쪽 항에서 (m - i)는 nums1 배열의 오른쪽 부분의 개수이고 (n - j)는 nums2 배열의 오른쪽 부분의 개수이다.

 

그리고 [식2]와 [식3]도 마찬가지인 맥락이고 +1이 오른쪽 항에 붙었는지와 왼쪽 항에 붙었는지의 차이일 뿐이다.

 

[식1]은 두 배열 개수 총 합(m + n)이 짝수일 때이다.

 

다시 말해 왼쪽 부분의 개수와 오른쪽 부분의 개수가 정확히 같은 경우는 요소들 총 개수가 짝수일 때 가능하다.

 

[식2]는 왼쪽 부분에 1개가 더 많은 것을 나타내고 [식3]은 오른쪽 부분에 1개가 더 많은 경우를 나타낸다.

(+1을 붙인 쪽이 더 적기 때문에 +1을 붙여주는 것이다.)

 

그리고 [식2]와 [식3]은 요소들 총 개수가 홀수일 때이다.

 

위 식들을 j 기준으로 정리하면 아래와 같다.

 

$$ [식1]　j = (m + n) / 2 - i $$

$$ [식2]　j = (m + n + 1) / 2 - i $$

$$ [식3]　j = (m + n - 1) / 2 - i $$

 

여기서 바로 알 수 있는 것은 우리가 첫 번째 조건을 만족시키기 위해 i를 선택하면 j는 결정된다는 것이다.

 

즉 i를 선택하고 위 식으로 j를 구하면 첫 번째 조건을 만족시킨다는 것이다.

 

이제 위에서 어떤 식을 써야 할까?

 

두 배열의 개수가 각각 짝수일 때와 홀수일 때 경우의 수를 다 구분하고 위 식 중에서 골라서 사용해야 할까?

 

모든 언어가 그런 것은 아니지만 C/C++ 기준으로 저 식을 계산할 때 소수점은 버려지게 된다.

(m, n, i, j의 타입이 정수일 때)

 

그래서 [식2]는 [식1]을 사용해야 할 경우도 포함하게 된다.

([식2]에서 m + n이 짝수면 +1로 인해 분자가 홀수가 되고 2로 나누면 ~.5가 되고 이것은 버려지게 된다.

결국 [식1]에서 구하는 j와 같게 된다.)

 

[식3]은 [식1]을 포함하지 못하기 때문에 [식3]으로 모든 경우를 다룰 수는 없다.

 

그래서 우리는 [식2]를 사용하여 m과 n의 짝수와 홀수 구분 없이 사용하기로 한다.

 

이제 두 번째 조건을 식으로 표현하면 아래와 같다.

 

$$ nums1[i - 1] <= nums2[j] && nums2[j - 1] <= nums1[i] $$

 

nums1 배열과 nums2 배열은 이미 정렬되어 있기 때문에 두 번째 조건을 4개 요소만 사용해서 검사할 수 있다.

 

이제 두 조건 모두 만족하는 i와 j를 찾았고 어느 요소들이 중간에 위치할 수 있는지 알게 됐다고 하자.

 

m + n이 짝수면 중간에 있는 두 요소들의 중간값을 계산하면 된다.

 

m + n이 홀수면 딱 중간에 있는 요소의 값이 답이다.

 

먼저 m + n이 짝수일 때 중간에 위치할 두 요소를 찾아야 한다.

 

위의 4개 요소 중에서 쉽게 찾을 수 있다.

 

nums1[i - 1]과 nums2[j - 1] 중에 하나가 왼쪽 부분의 최댓값이 될 것이고

 

nums1[i]와 nums2[j] 중에 하나가 오른쪽 부분의 최솟값이 될 것이다.

 

왼쪽 부분의 최댓값과 오른쪽 부분의 최솟값으로 중간값을 계산하면 된다.

 

다음으로 m + n이 홀수일 때이다.

 

짝수일 때와 마찬가지로 왼쪽 부분의 최댓값과 오른쪽 부분의 최솟값을 구한다.

 

여기서 둘 중 하나를 결정하면 되는데 우리가 첫 번째 조건의 [식2]를 사용할 경우를 생각해 보자.

 

[식2]는 왼쪽 부분에 요소가 1개 더 많은 경우를 나타낸다고 했다.

 

왼쪽 부분에 1개 더 많으니 왼쪽 최댓값과 오른쪽 최솟값 중에 왼쪽 최댓값이 딱 중간에 있는 것이므로 왼쪽 최댓값을 선택하면 된다.

 

마지막으로 i를 선택하는 방법을 생각해 보자.

 

그냥 0부터 m - 1까지 순서대로 조건들을 검사할 수도 있다.

 

하지만 두 번째 조건을 검사하고 다음으로 검사해 볼 요소가 선택했던 요소 기준으로 왼쪽에 있는지 오른쪽에 있는지 알 수 있다.

 

nums1[i]가 nums2[j - 1]보다 작으면 다음으로 선택할 요소를 i 이후에 있는 요소로 선택해야 한다.

 

nums2[j]가 nums1[i - 1]보다 작으면 다음으로 선택할 요소를 i 이전에 있는 요소로 선택해야 한다.

 

앞서 잠깐 언급했던 이진 탐색을 사용하면 보다 빠르게 찾을 수 있다.

 

이제 이 내용들을 정리하여 코드로 표현하면 아래와 같다.

 

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        if (nums1.size() > nums2.size())
        {
            nums1.swap(nums2);
        }

        const int n = nums1.size();
        const int m = nums2.size();
        int left = 0, right = n;
        while (left <= right)
        {
            int i = (right - left + 1) / 2 + left;
            int j = (n + m + 1) / 2 - i;
            int nums1_1 = i > 0 ? nums1[i - 1] : INT_MIN;
            int nums1_2 = i < n ? nums1[i] : INT_MAX;
            int nums2_1 = j > 0 ? nums2[j - 1] : INT_MIN;
            int nums2_2 = j < m ? nums2[j] : INT_MAX;
            if (nums1_1 > nums2_2)
            {
                right = i - 1;
            }
            else if (nums2_1 > nums1_2)
            {
                left = i + 1;
            }
            else
            {
                int leftMax = max(nums1_1, nums2_1);
                int rightMin = min(nums1_2, nums2_2);
                if (0 != (n + m) % 2)
                {
                    return leftMax;
                }
                else
                {
                    return (leftMax + rightMin) / 2.0;
                }
            }
        }

        return 0.0;
    }
};

 

m <= n을 보장하기 위해서 처음에 swap하는 부분이 있다.

 

m <= n이 보장되어야 nums2 배열에서 out of bounds 문제가 발생하지 않는다.

 

그리고 i와 j가 배열의 경계 부분에 있을 때를 위해 INT_MIN과 INT_MAX를 대신 사용한 부분도 있다.