[leetcode 매일 풀기] 4. median of two sorted arrays hard

한동안 일반적인 개발을 하다 보니, 뭔가 날카로움(엣지케이스에 대한 대응 능력, 코드를 깔끔하게 짜는 능력)이 부족해지는 느낌을 받았습니다.

 

그래서 시작하기로 마음먹은 leetcode hard mode 문제 풀기 챌린지입니다.

 

2주정도 전에 시작했는데, 푼시간과 올리는 시간을 텀을 조금 주고 하나씩 올리면서 복기겸, 코멘트를 남겨볼까 합니다.

 

https://leetcode.com/problems/median-of-two-sorted-arrays/description/

 

이 문제는 두개의 정렬된 list가 주어졌을 때 중간에 있는 list를 구하는 알고리즘을 짜는 문제입니다.

 

오랜만에 알고리즘을 구현하려니, 엣지케이스가 많은 경우에 상당히 고생을 많이하면서. 이런 부분에 있어서 많이 무뎌졌구나하는 느낌을 다시금 받았습니다.

 

문제의 제약조건이 time complexity O(n+m) 안에 구현하라는 것이어서, 대충 binary search를 하면 되겠지라고 생각했는데.

median이라는 조건이 edge case를 더욱 어렵게 만들었던 문제입니다.

 

결국 solution을 보고, 가다듬어서 제출을 한 첫번재 hard 문제 solution을 공유합니다.

from typing import List


class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        n1 = len(nums1)
        n2 = len(nums2)

        if n1 > n2:
            return self.findMedianSortedArrays(nums2, nums1)

        n = n1 + n2
        num_left = (n + 1) // 2

        # binary search in nums1
        left = 0
        right = n1
        while left <= right:
            mid1 = (left + right) // 2
            mid2 = num_left - mid1

            l1 = float('-inf')
            l2 = float('-inf')
            r1 = float('inf')
            r2 = float('inf')

            if mid1 < n1:
                r1 = nums1[mid1]
            if mid2 < n2:
                r2 = nums2[mid2]
            if mid1 - 1 >= 0:
                l1 = nums1[mid1 - 1]
            if mid2 - 1 >= 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.0
            elif l1 > r2:
                right = mid1 - 1
            elif l2 > r1:
                left = mid1 + 1


assert Solution().findMedianSortedArrays([1, 3], [2]) == 2.0
assert Solution().findMedianSortedArrays([1, 2, 3], [-3, -2, -1, 0]) == 0.0
assert Solution().findMedianSortedArrays([1, 2], [3, 4]) == 2.5