[코딩 테스트] 리트코드 4 - Median of Two Sorted Arrays (Hard) in Python

Given two sorted arrays nums1 and nums2 of size m and n respectively, return the median of the two sorted arrays.

The overall run time complexity should be O(log (m+n)).

 

두 정렬된 배열들 num1와 nums2가 주어질 때, 둘의 median 값을 리턴하라.

 

Example 1:

Input: nums1 = [1,3], nums2 = [2]

Output: 2.00000

Explanation: merged array = [1,2,3] and median is 2.

 

Example 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.

 

Example 3:

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

Output: 0.00000

 

풀이

(1) 두 배열을 합친 후 정렬 - 무작정 풀기 [Brute Force]

두 배열이 합쳐진 (m+n) 후 정렬하면 시간 복잡도는 O(log(m+n)), 문제의 요구와 맞는다.

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        nums = nums1 + nums2
        nums.sort()
        
        n = len(nums)
        if n % 2 == 1: #합쳐진 배열의 원소 갯수가 홀수개일 땐 가운데 값 리턴 
            answer = nums[n // 2]
        else: # 합쳐진 배열의 원소 갯수가 짝수개일 땐 가운데 두 값의 평균값 리턴
            answer = (nums[n // 2 - 1] + nums[n // 2]) / 2.0
            
        return answer

 

(2) Median의 정의 이용

출처 - https://www.youtube.com/watch?v=LPFhl65R7ww 

Median의 정의

: 어떠한 수의 집합을 (1) 같은 크기의 둘로 나누되,

(2) 한쪽은 median보다 작거나 같게, 다른 한 쪽은 median보다 크거나 같게 나누는 방법이다.

 

예시

nums1에는 6개의 원소가,

nums2에는 8개의 원소가 있다고 치자.

합쳐진 배열에서, nums1에서는 2개의 원소를 '왼쪽 그룹'에 넣는다면,

nums2에서는 5개의 원소를 '왼쪽 그룹'에 넣어줘야

'왼쪽 그룹'의 원소의 갯수(2 + 5)와 '오른쪽 그룹'의 원소의 갯수(4 + 3)가 같게 될 것이다. [Median 조건 (1)]

 

그리고 Median의 조건 (2)에 의해, '왼쪽 그룹'의 모든 원소는 '오른쪽 그룹'의 모든 원소의 값보다 작거나 같아야 한다.

이 때, nums1와 nums2는 이미 정렬되어 있으므로,

x2 < x3

y5 < y6

인 건 당연하다.

따라서 중요한 건

x2 <= y6

y5 <= x3

인 두 그룹으로 나눠야 한다는 것이다.

 

이 때 Median은

max(x2, y5)와 min(y6, x3)의 평균값

이 될 것이다!

 

만약 합쳐진 배열의 원소 갯수가 홀수개라면,

'왼쪽 그룹'의 원소 갯수가 '오른쪽 그룹'의 원소 갯수보다 하나가 더 많을 것이고,

이 때의 Median은

max(x2, y5)

가 된다.

 

 

위와 같은 두 그룹을 Binary Search (이진탐색)을 통해 구하는 것이 문제의 핵심이 되겠다.

 

이진탐색의 알고리즘

Pseudo-Code

# '왼쪽 그룹'의 원소의 길이는 다음과 같이 구한다.
# partitionLength = (len(nums1) + len(nums2) + 1) // 2

while 찾을 때까지:
    binary_search_index = (start + end) / 2
    index_for_'왼쪽그룹'_of_nums1 = binary_search_index
    index_for_'왼쪽그룹'_of_nums2 = partitionLength - index_for_'왼쪽그룹'_of_nums1
    
    Found: # 답을 찾은 경우는 아래와 같은 경우
    	maxLeftX <= minRightY # nums1의 왼쪽 그룹 중 max값이 nums2의 오른쪽 그룹 중 min값보다 작거나 같음
        maxLeftY <= minRightX
    elif maxLeftX > minRightY: # nums1의 왼쪽 그룹의 max값이 너무 큰 경우다. 따라서 왼쪽 그룹을 줄여줌
    	move towards left in nums1
    else:
    	move towards right in nums1

 

 

 

이진탐색 이용 파이썬 코드

위의 이론들을 조합하면 파이썬 코드로는 다음과 같이 짤 수 있다.

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        if len(nums1) > len(nums2):
            return self.findMedianSortedArrays(nums2, nums1)
        
        partitionLength = (len(nums1)+len(nums2)+1)//2
        isEven = (len(nums1)+len(nums2))%2 == 0
        
        start = 0
        end = len(nums1)
        while start <= end:
            nums1Idx = (start+end)//2
            nums2Idx = (partitionLength - nums1Idx)
            
            # Edge Cases - 왼쪽 그룹이 비었으면 -INF값을,
            # 오른쪽 그룹이 비었으면 +INF값을 적용
            maxLeft1 = nums1[nums1Idx-1] if nums1Idx != 0 else -987654321
            minRight1 = nums1[nums1Idx] if nums1Idx != len(nums1) else 987654321
            
            maxLeft2 = nums2[nums2Idx-1] if nums2Idx != 0 else -987654321
            minRight2 = nums2[nums2Idx] if nums2Idx != len(nums2) else 987654321
            
            if maxLeft1 <= minRight2 and maxLeft2 <=  minRight1:
                if isEven:
                    return (max(maxLeft1,maxLeft2)+min(minRight1,minRight2))/2
                else:
                    return max(maxLeft1, maxLeft2)
            elif maxLeft1 > minRight2:
                end = nums1Idx - 1
            else:
                start = nums1Idx+1
            

 

이 때 시간복잡도는 O(log(min(m,n))으로 더욱 빠른 것을 볼 수 있다.

 

새로운 배열을 만들 필요 없이 원래의 배열들에서 처리가 가능하니 공간복잡도 면에서도 첫 번째 문제보다 효율적이다.