[Leetcode] 4. Median of Two Sorted Arrays

문제 바로가기

O(m+n) 풀이

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        len1, len2 = len(nums1), len(nums2)
        isEven = (len1 + len2) % 2 == 0
        idx1, idx2 = 0, 0
        lastNum = 0
        median = 0
        
        for i in range(int((len1 + len2)/2)):
            if idx1 == len1:
                lastNum = nums2[idx2]
                idx2 += 1
            elif idx2 == len2:
                lastNum = nums1[idx1]
                idx1 += 1
            elif nums1[idx1] <= nums2[idx2]:
                lastNum = nums1[idx1]
                idx1 += 1
            else:
                lastNum = nums2[idx2]
                idx2 += 1

        # 중간값 도출
        if idx1 == len1:
            median = nums2[idx2]
        elif idx2 == len2:
            median = nums1[idx1]
        elif nums1[idx1] <= nums2[idx2]:
            median = nums1[idx1]
        else:
            median = nums2[idx2]

        if isEven:
            median = (median + lastNum) / 2

        return median

O(log(m+n)) 풀이

binary search처럼 반절씩 줄여나가며 탐색

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        len1, len2 = len(nums1), len(nums2)
        # 길이가 더 짧은 리스트가 nums1이 되게 한다
        if len1 > len2:
            nums1, nums2, len1, len2 = nums2, nums1, len2, len1
        lenSum = len1 + len2
        halfSum = lenSum // 2
        rear1, rear2 = 0, 0
        rearMin, rearMax = 0, len1

        while rearMin <= rearMax:
            # 두 리스트의 처음부터 rear1-1, rear2-1까지의 요소는 half 개이다. 
            rear1 = (rearMin + rearMax) // 2   # nums1의 index
            rear2 = halfSum - rear1       # nums2의 index
            
            if rear1 < len1 and nums1[rear1] < nums2[rear2-1]:
                # rear1이 너무 작음
                rearMin = rear1 + 1
            elif rear1 > 0 and nums1[rear1-1] > nums2[rear2]:
                # rear1이 너무 큼
                rearMax = rear1 - 1
            else:
                # 완성
                if rear1 == len1:
                    median = nums2[rear2]
                elif rear2 == len2:
                    median = nums1[rear1]
                else:
                    median = min(nums1[rear1], nums2[rear2])
                # 짝수
                if lenSum % 2 == 0:
                    if rear1 == 0:
                        median += nums2[rear2-1]
                    elif rear2 == 0:
                        median += nums1[rear1-1]
                    else:
                        median += max(nums1[rear1-1], nums2[rear2-1])
                    median /= 2
                return median