Leetcode 4. Median of Two Sorted Arrays with 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)).

두 정렬된 배열 nums1와 nums2가 주러질 때 이 두 배열을 합친 배열의 중간값을 리턴하라.

예시

Input: nums1 = [1,3], nums2 = [2]
Output: 2.00000
Explanation: 합쳐진 배열인 [1,2,3] 의 중간값은 2이다.

제한

• nums1.length == m
• nums2.length == n
• 0 <= m <= 1000
• 0 <= n <= 1000
• 1 <= m + n <= 2000
• 10^6 <= nums1[i], nums2[i] <= 10^6

풀이법

1. TIME COMPLEXITY O(nlogn)

두 배열을 무식하게 합치고 정렬한 뒤 중앙값을 찾는다.

배열 정렬의 시간 복잡도가 nlogn이기에 이 또한 nlogn이 된다.

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        nums3 = nums1 + nums2
        nums3.sort()
        if len(nums3) % 2 == 1:
            return nums3[len(nums3) // 2]
        else:
            return (nums3[len(nums3) // 2 - 1] + nums3[len(nums3) // 2]) / 2
         

2. TIME COMPLEXITY O(n)

두 배열이 이미 정렬되어 있다는 사실을 이용하여

두 배열을 앞에서부터 끝까지 순회하여 정렬된 합쳐진 배열을 만든 뒤 중앙값을 찾는다.

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        size = len(nums1) + len(nums2)
        idx1 = idx2 = 0
        res = []
        for i in range(size):
            if idx1 < len(nums1) and idx2 < len(nums2):
                if nums1[idx1] < nums2[idx2]:
                    res.append(nums1[idx1])
                    idx1 += 1
                else:
                    res.append(nums2[idx2])
                    idx2 += 1
            elif idx1 < len(nums1):
                res.append(nums1[idx1])
                idx1 += 1
            else:
                res.append(nums2[idx2])
                idx2 += 1
        if len(res) % 2 == 1:
            return res[len(res) // 2]
        else:
            return (res[len(res) // 2 - 1] + res[len(res) // 2]) / 2
            

3. Quickselect

전능하신 ChatGPT에게 물어본 결과 Quickselect를 쓰면 된다고 하여서

Quickselect를 사용해 보았다.

두 배열을 더해 Quickselect를 돌리면 된다.

def quick(lst, start, end):
    pivotValue = lst[start]
    i = start + 1
    j = end
    while i <= j:
        while i <= end and lst[i] < pivotValue:
            i += 1
        while j > 0 and lst[j] > pivotValue:
            j -= 1
        if i <= j:
            lst[i], lst[j] = lst[j], lst[i]
            i, j = i+1, j-1
    lst[start], lst[j] = lst[j], lst[start]
    return j

def findIndex(lst, index):
    start = 0
    end = len(lst)-1
    while start <= end:
        p = quick(lst, start, end)
        if p == index: return lst[p]
        elif p > index:
            end = p - 1
        else:
            start = p + 1


class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        nums = nums1 + nums2

        if len(nums) % 2 == 0:
            return (findIndex(nums, len(nums) // 2 - 1) + findIndex(nums, len(nums) // 2)) / 2
        else:
            return findIndex(nums, len(nums) // 2)

근데 성능이 오히려 별로가 되었다..