[python] LeetCode 4번: Median of Two Sorted Arrays(Hard)

27 Nov 2020 in Algorithm on LeetCode

---

Problem

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

nums1.length == m
nums2.length == n
0 <= m <= 1000
0 <= n <= 1000
1 <= m + n <= 2000
-10⁶ <= nums1[i], nums2[i] <= 10⁶

두 정렬된 배열 nums1, nums2가 각각 m, n의 길이로 주어졌을 때, 두 배열을 합친 배열의 중앙값을 구하여라.

문제 링크 : 리트코드 4번

---

<Try>

두 배열이 정렬되어 있으므로 두 배열의 가장 앞에 있는 숫자를 비교하여 더 작은 수를 차례대로 나열하고 중앙값에 해당되는 인덱스에 도달하면 값을 반환하는 방식을 사용했다.

두 배열 길이의 합이 짝수이면 중앙값은 가운데 2개의 평균을 반환하고 홀수이면 가운데 1개의 값만 반환하면 된다는 것을 고려해야 한다.

---

<Solution>

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m = len(nums1)
        n = len(nums2)
        merged_array = []
        # p1, p2는 각각 nums1, nums2의 인덱스를 가리키는 포인터이다.
        p1 = 0
        p2 = 0
        # 도달해야할 중앙값의 인덱스이다.
        pivot = (m+n)//2 + 1
        # merged_array에 나열된 숫자의 개수
        curr_idx = 1
        while True:
            # 제한조건에 10^6보다 클 수 없기 때문에 배열이 끝나면 10^6+1 대입
            v1 = nums1[p1] if p1 < m else 1000001
            v2 = nums2[p2] if p2 < n else 1000001
            if v1 < v2:
                p1 += 1
                merged_array.append(v1)
            else:
                p2 += 1
                merged_array.append(v2)
            
            if curr_idx == pivot:
                if (m+n) % 2 == 0:  # 총 배열의 길이가 짝수
                    return (merged_array[-1] + merged_array[-2]) / 2
                
                else:   # 홀수
                    return merged_array[-1]
            
            curr_idx += 1
                

두 배열길이의 합 m+n의 절반을 순회한다. n>=m이라 가정하면 시간 복잡도는 O(n)이다.