Leetcode 4 Solution

Leetcode Question Link

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

Thinking Process

This is a hard question. There are a couple of steps in the thinking process.

Step 1: Set The Time Complexity Goal

The question is trivial if our goal is O(m + n). Just use two pointers, and walk through the two arrays. Always walk the pointer who’s pointing to a smaller value. Stop once we walk half of the elements.

Apparently, there has to be something better. And the first thing we should think of is that if linear solution is not good enough, we should try to search for log solution. So our time complexity goal should be O(lg(m + n)).

Step 2: Transformation

Before we tackle the question, let’s transform this question into another one, which is to get the kth largest element of two arrays. If m + n is even, k = (m + n) // 2. Otherwise, k = (m + n) // 2 + 1. Apparently, the kth largest element is an easier question to tackle.

Now the question becomes how do we get the kth largest element of the two arrays within O(m + n) time.

Step 3: Derive The Algorithm

Whenever we want to design a log algorithm, we almost always need to think of a way to remove a part of the elements in each iteration. And that number of elements to be removed has to be propotional to the current scale. For example, in binary search, each iteration, we eliminate half of the elements. In this case, let’s try to remove k // 2 elements from the two arrays in each iteration.

Since we are going to remove elements from the front, let’s assume the answer of the question is F(i, j, k) after we removed nums1[:i] and nums2[:j]. In other words, F(i, j, k) is the kth largest element of array nums1[i:] and nums2[j:]. Now the question is to calculate F(i, j, k).

First of all, if i == m or j == n, it means we just need to get the kth element from the other array.

Second, if k == 1, we’ll just grab the smaller element of nums1[i] or nums2[j].

Otherwise, let’s try to walk k // 2 steps from each array. Note that if we exceed the end of an array, stop at the last element. In other words,

next_i = min(len(nums1) - 1, i + k // 2 - 1)
next_j = min(len(nums2) - 1, j + k // 2 - 1)

Now, there are two cases:

• If nums1[next_i] is smaller than nums2[next_j], it means that none of the elements from i to next_i in nums1 could be the target, so the answer should be F(next_i + 1, j, k - (next_i + 1 - i)).
• If nums1[next_i] is larger than nums2[next_j], it means that none of the elements from j to next_j in nums2 could be the target, so the answer should be F(next_i, j + 1, k - (next_j + 1 - j)).

Till now, we have derived the solution. Each iteration, we remove k // 2 elements, so the algorithm can finish within O(lg(k)) time.

NOTE: Although our thinking process is recursive, the implementation doesn’t have to be so, because each time we are only visiting one branch, just like binary search.

Time & Space Complexity

• Time Complexity: O(lg(m + n))
• Space Complexity: O(1)

Solution

class Solution:
    def findKthElement(self, nums1, nums2, k):
        i = 0
        j = 0

        while True:
            if i == len(nums1):
                return nums2[j + k - 1]
            elif j == len(nums2):
                return nums1[i + k - 1]

            if k == 1:
                return min(nums1[i], nums2[j])

            next_i = min(len(nums1) - 1, i + k // 2 - 1)
            next_j = min(len(nums2) - 1, j + k // 2 - 1)

            if nums1[next_i] < nums2[next_j]:
                k -= next_i + 1 - i
                i = next_i + 1
            else:
                k -= next_j + 1 - j
                j = next_j + 1

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

        if (n1 + n2) % 2 == 0:
            return ((self.findKthElement(nums1, nums2, (n1 + n2) // 2) + self.findKthElement(nums1, nums2, (n1 + n2) // 2 + 1)) / 2)
        else:
            return self.findKthElement(nums1, nums2, (n1 + n2) // 2 + 1)

Learning

• Set a time complexity goal if possible.
• Always try to transform the question to a simple one if possible in the first place.