[Python 3] LEETCODE 4 Median of Two Sorted Arrays

최근, LEETCODE라는 웹 사이트에서 알고리즘 공부를 하고 있다.

백준 온라인 저지는 조금 "구현"적인 느낌이 강하고, 문제에 미사여구 (재미를 주는 부분) 또한 많다. 한국어로 이루어져 있기에 조금 더 문제 읽기가 좋다.

그러나 Leetcode에서는 이 보다는 알고리즘 그 자체를 더욱 강하게 다루는 느낌이다. 같은 알고리즘이라고 하더라도, 분명히 문제가 주는 느낌이 다르다. 결코 같은 문제를 푼다는 느낌이 안 든다. 조금 더 코딩 인터뷰나, 이해도를 키우기 위해서는 반드시 Leetcode도 풀어봐야겠다, 라는 생각을 했다.

아직까지 풀었던 Leetcode 문제는 Easy나 Medium 난이도였는데, 처음으로 Hard 난이도를 자랑하는 문제를 마주쳤다.

번역을 해보자면,

2개의 정렬된 배열 nums1과 nums2가 각자 m과 n의 사이즈로 주어져 있을 때, 두 개의 정렬된 배열의 중앙값을 return 하여라.
전체적인 실행 시간의 복잡도는 O(log(m+n))이여야 한다.

 

처음 이 문제를 풀이할 때, 투 포인터를 통해서 풀려고 했다. 각각 nums1와 nums2의 인덱스인 i,j를 각각 설정하여, 두 배열의 크기의 절반 (짝수인 경우에는 중앙에 제일 가까운 두 수의 평균)을 구한다는 생각이였다.

하지만, 이렇게 풀게 될 경우 ((nums1의 길이)+(nums2의 길이))의 절반만큼의 시간 복잡도, 즉 O(n)의 시간 복잡도가 걸린다. 이렇게 되는 경우 문제 조건에도 위배되고, 시간 초과를 겪을 것이다.

 

이런 Linear한 리스트에서 보통 log 시간을 사용하는 대표적인 방법은 Binary Search이다. 특히, 배열이 이미 나열되어 있는 상태이기 때문에, 더욱 적합하다. 하지만, 두 개의 배열이 동시에 있기 때문에, 도대체 어떤 방법으로 해결해야할지 큰 고민이 되었다. 결국 다른 사람들의 내 코드를 섞었다.

 

class Solution:
    def findMedianSortedArrays(self, nums1, nums2):
        n1 = len(nums1)
        n2 = len(nums2)
        if n1 > n2:
            return self.findMedianSortedArrays(nums2, nums1)
        n = n1 + n2
        l = (n1 + n2 + 1) // 2 
        low = 0
        high = n1        
        while low <= high:
            mid1 = (low + high) // 2 
            mid2 = l - mid1 
            l1 = -1e10
            l2 = -1e10
            r1 = 1e10
            r2 = 1e10
            if mid1 < n1:
                r1 = nums1[mid1]
            if mid2 < n2:
                r2 = nums2[mid2]
            if mid1 > 0:
                l1 = nums1[mid1-1]
            if mid2 > 0:
                l2 = nums2[mid2-1]
            if l1 <= r2 and l2 <= r1:
                if n % 2 == 1:
                    return max(l1, l2)
                else:
                    return (max(l1, l2) + min(r1, r2)) / 2
            elif l1 > r2:
                high = mid1 - 1
            else:
                low = mid1 + 1

내가 제출한 코드이다. 

 

nums1과 nums2 중에서 무조건적으로 nums1의 길이가 작게 해야 코드가 간결해질 것 같아, nums2의 길이가 nums1의 길이보다 더 길 경우 findMedianSortedArray를 다시 call 해주며 nums1의 크기가 nums2의 크기보다 더 작도록 하였다.

여기서 n은 전체 코드의 길이, l은 절반(중앙값이 있어야 할 자리)로 설정하였고, 이분 탐색은 nums1의 첫번쩨 인덱스부터 마지막 인덱스까지 실시하였다. nums1의 이분탐색을 진행하여 보자.

 

이분 탐색 안을 살펴보자. mid1은 이분탐색을 진행하는 부분의 index이다. mid2는 l-mid1으로 설정하였다.

이는 왜냐하면, "중앙값의 후보가 mid1이면, 다른 중앙값의 후보는 mid2일 수 밖에 없다. 두 수가 l로 더해지지 않는다면, 두 수 중 정답이 있더라도, 정말 이 수가 중앙값인지 확인 할 수 있는 방법이 없기 때문이다."

여기서 l1, r1은 nums1에서 mid1을 기준으로 좌우, l2,r2는 nums2에서 mid2를 기준으로 좌우로 확인하는 경우이다. 여기서 out of bounds 일수도 있으므로, 초기값을 각각 -1e10과 1e10 (마이너스 무한대와 무한대)로 설정하였다.

 

만약 l1<=r2, l2<=r1, 즉 서로의 좌우 값이 대소를 만족한다면, 결과를 return해준다.

당연히 중앙값이기 때문에 숫자가 짝수개 있으면 두 수의 평균을 출력해주어야 한다.

 

앞으로 Leetcode도 많이 업로드해볼 생각이다.