LeetCode - 4. Median of Two Sorted Arrays

문제 개요

난이도 - HARD 사용 언어 - C++

정수 값이 담겨있으며 정렬된 두 배열이 넘어올 때 두 배열의 중앙값(Median)을 계산하는 문제입니다.

중앙값은 정렬된 수의 배열 중 중앙에 위치하는 값을 말합니다. 수의 개수가 홀수라면 하나가 선택되고, 짝수라면 두 수를 골라 평균을 계산해 중앙값을 구합니다.

input example

위 이미지와 같이 정렬된 정수 배열 2개가 전달된다면, 두 배열의 수를 나란히 나열했을 때 1, 1, 2, 3, 10, 10, 20, 45, 55, 70이 되며, 10, 10이 중간값으로 선택되니 (10+10)/2 = 10, 즉 최종적으로 10이 중간값이 됩니다.

Reference Video

이 문제를 해결하기 위해 위 영상의 설명을 참고했습니다.

문제 - LeetCode 4. Add Two Numbers

풀이

Solution

중간값을 구하기 전에 먼저 문제의 제한사항을 확인했습니다.

소스 코드의 시간 복잡도는

$O(log(m+n))$

(m, n은 각 배열의 길이)을 넘지 않아야 합니다. 이 때문에 두 배열을 하나로 합쳐서 중간값을 구할 수는 없습니다.

두 배열을 하나로 합치려면 m, n을 순회해야 하므로

$O(m+n)$

의 시간 복잡도가 발생하는데, 이는

$O(m+n) > O(log(m+n))$

이므로 배열을 합치지 않고 중간값을 찾아야 합니다.

이를 해결하기 위해 이진 탐색과 비슷한 논리를 이용하는 방법을 떠올렸습니다.

이진 탐색을 수행하려면 배열이 정렬되어 있어야 하며, 맨 처음 중간값을 골랐을 때 정렬된 배열이라면 중간값의 왼쪽 값은 당연히 중간값보다 작거나 같아야 하고, 오른쪽 값은 무조건 중간값보다 크거나 같아야 합니다.

median example

두 번째 특성은, 중간값을 선택하면 중간값보다 작은 값의 개수와 큰 값의 개수가 동일하다는 것입니다.

이러한 특성을 이용해 두 배열을 합치지 않고 중앙값을 구할 수 있습니다.

배열의 적절한 Left, Middle, Right 범위 구하기

먼저 배열의 중앙값을 구하기 위해, 중앙값을 골랐을 때 작은 값과 큰 값의 개수가 같다는 성질을 이용해 적절한 Left(작은 값의 범위), Middle(중앙값), Right(큰 값의 범위)를 구합니다.

median example 2

두 배열의 길이는 10이므로, 두 배열을 하나로 합쳤다고 생각하면 5개씩 두 그룹으로 나눌 수 있습니다. 이때 중앙값을 기준으로 Left에 5개, Right에 5개의 값이 위치합니다.

먼저 A 배열의 중앙값(A_left)을 선택합니다. 위 그림에서는 3번째 값이 선택되었고, Left에 위치할 값 5개를 채우기 위해 B 배열에서 2번째 값을 골라 B의 중앙값(B_left)으로 선택합니다. 이렇게 총 5개의 Left가 선택되면 나머지 5개는 자동으로 Left보다 큰 값, 즉 Right가 됩니다.

이제 Left, Right로 나눈 결과가 적절한지 확인합니다. Left의 값들은 무조건 Right의 값보다 작아야 합니다.

1, 1, 1, 2, 3 <-Left Right-> 2, 10, 45, 55, 70

Left에 3이라는 값이 있지만 Right에는 3보다 작은 2가 포함되어 있어, 지금 선택된 Left와 Right는 적절하지 않음을 알 수 있습니다.

median example 4

사람의 눈으로 보면 바로 구분이 되지만, 이를 확실하게 판별하는 조건은 A_left(A의 중앙값)가 B_left+1(B의 중앙값의 다음 값)보다 작아야 한다는 것입니다. 반대로 B_left <= A_left+1도 성립해야 합니다.

이를 해결하기 위해 A의 중앙값 index를 1 감소시키고, B의 중앙값을 1 증가시킵니다.

median example 3

다시 적절한 분류인지 확인합니다.

1, 1, 1, 2, 2 <-Left Right-> 3, 10, 45, 55, 70

확인해보니 적절하게 분류되었습니다. 이제 중앙값을 선택해야 하는데, 위처럼 정리한 것은 편의상 일렬로 나란히 나열한 것일 뿐, 실제로는 두 배열의 index를 각각 가지고 있기 때문에 이를 고려해 적절히 선택할 방법이 필요합니다.

이는 간단하게 수행할 수 있습니다. Left로 분류된 값 중 가장 큰 값과 Right로 분류된 값 중 가장 작은 값이 중앙값 후보가 됩니다.

즉 (max(A_left, B_left) + max(A_right, B_right))/2가 중앙값이 됩니다.

제출 결과

Solution 1 result

실행 시간은 20ms로, 다른 C++ 제출자에 비해 96%가량 좋은 성능을 보이는 코드를 작성할 수 있었습니다.

코드 전문

#include <vector>
#include <algorithm>
#include <iostream>
#include <limits>

class Solution 
{
public:
    double findMedianSortedArrays(std::vector<int>& nums1, std::vector<int>& nums2) 
    {
        if (nums1.size() == 0 && nums2.size() == 0)
        {
            return 0.0;
        }

        int totalLength = nums1.size() + nums2.size();
        int half = totalLength / 2;
        
        std::vector<int>& A = nums1;
        std::vector<int>& B = nums2;

        if (nums1.size() < nums2.size())
        {
            auto tmp = A;
            A = nums2;
            B = tmp;
        }

        if (A.size() == 0)
        {
            int mid = B.size() / 2;
            if ((B.size() % 2) == 0)
            {
                return (B[mid - 1] + B[mid]) / 2.0;
            }
            else
            {
                return (double)B[mid];
            }
        }
        else if (B.size() == 0)
        {
            int mid = A.size() / 2;
            if ((A.size() % 2) == 0)
            {
                return (A[mid - 1] + A[mid]) / 2.0;
            }
            else
            {
                return (double)A[mid];
            }
        }

        int l = 0;
        int r = A.size() - 1;

        do 
        {
            int aLeftMidIndex = (l + r) /2;
            int bLeftMidIndex = half - aLeftMidIndex - 2;

            int aLeft = aLeftMidIndex >= 0 ? A[aLeftMidIndex] : std::numeric_limits<int>::min();
            int aRight = (aLeftMidIndex + 1) < A.size() ? A[(aLeftMidIndex + 1)] : std::numeric_limits<int>::max();
            int bLeft = bLeftMidIndex >= 0 ? B[bLeftMidIndex] : std::numeric_limits<int>::min();
            int bRight = (bLeftMidIndex + 1) < B.size() ? B[(bLeftMidIndex + 1)] : std::numeric_limits<int>::max();

            if (aLeft <= bRight && bLeft <= aRight)
            {
                if ((totalLength % 2) == 0)
                {
                    return (std::max(aLeft, bLeft) + std::min(aRight, bRight)) / 2.0;
                }
                else
                {
                    return (double)std::min(aRight, bRight);
                }
            }
            else if (aLeft > bRight)
            {
                r--;
            }
            else if (bLeft > aRight)
            {
                r++;
            }
        } while(true);
    }
};