[LeetCode] Median of Two Sorted Arrays - Python 풀이

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

두 개의 정렬된 List 합쳤을 때 중앙값을 구하는 문제이다. 

단, 문제에서 O(log(m+n)) 시간 내에 풀이하라고 주어졌다. 

 

1. Merge Sort 알고리즘처럼 하나씩 대소비교를 하여 정렬하는 풀이 

• 시간복잡도 : O(m+n)
• Runtime : 90 ms
• Memory : 16.8MB

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        n = len(nums1)
        m = len(nums2)
        t = (n+m) # length of the merged array 

        left, right = 0, 0 # cursor for nums1, cursor for nums2
        merged_arr = []

        while left < n and right < m:
            if nums1[left] <= nums2[right]:
                merged_arr.append(nums1[left])
                left += 1
            else: 
                merged_arr.append(nums2[right])
                right += 1
        
        # 남은 숫자 
        if left < n: 
            merged_arr.extend(nums1[left:])
        if right < m:
            merged_arr.extend(nums2[right:])
            
        if len(merged_arr) % 2 == 0: # 짝수인 경우 
            mid = len(merged_arr)//2
            return (merged_arr[mid] + merged_arr[mid-1])/2 
        else: # 홀수인 경우 
            mid = len(merged_arr)//2 
            return merged_arr[mid]

 

2. Heap 정렬을 활용한 풀이 

• 시간복잡도 : O((m+n)*log(m+n))
• Runtime : 81ms
• Memory: 17MB

import heapq 

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        answer = 0 
        arr = nums1 + nums2 
        heapq.heapify(arr)
        
        t = len(arr)
        mid = t //2
        if t % 2 == 0: # 짝수 
            for i in range(t):
                now = heapq.heappop(arr)
                if i == mid-1:
                    answer += now
                if i == mid:
                    answer += now 
                    return answer /2
        else: # 홀수 
            for i in range(t):
                now = heapq.heappop(arr)
                if i == mid:
                    return now