Median of Two Sorted Arrays | Leetcode 4 | Optimal Binary Search Solution

The “Median of Two Sorted Arrays” problem is a classic and tricky coding interview question. It tests your understanding of merging, searching, and partitioning arrays.

In this blog, we’ll explain the problem, walk you through brute force, better, and optimal solutions, and make everything easy to understand with code comments, dry runs, and clear explanations.

Here’s the [Problem Link] to begin with.

Problem Statement

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

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.

Constraints:

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

Brute Force Solution (Merge and Find Median)

Intuition and Approach

How does this solution work?

1. Merging Two Sorted Arrays:
   Since both arrays are already sorted, we can use the two-pointer technique to merge them into a single sorted array. This is similar to the merge step in merge sort. We keep comparing the smallest elements from both arrays and add the smaller one to our result array.
2. Finding the Median:
   Once we have the merged array, finding the median is easy:
   • If the total number of elements is odd, the median is the middle element.
   • If the total number is even, the median is the average of the two middle elements.

Why does this work?
By merging, we get a single sorted array, and finding the median in a sorted array is straightforward.

When should you use this?
This approach is good for small arrays or when you want to build your understanding of the problem.

Code Implementation

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        result = []
        n, m = len(nums1), len(nums2)
        i, j = 0, 0
        while i < n and j < m:
            if nums1[i] <= nums2[j]:
                result.append(nums1[i])
                i += 1
            else:
                result.append(nums2[j])
                j += 1

        while i < n:
            result.append(nums1[i])
            i += 1
        while j < m:
            result.append(nums2[j])
            j += 1
        l = len(result)
        if l % 2 == 0:
            return (result[l // 2 - 1] + result[l // 2]) / 2
        return result[l // 2]

Better Solution (Find Median While Merging)

Intuition and Approach

How does this solution work?

1. Avoid Full Merge:
   Instead of merging the entire arrays, we only need to find the elements at the median positions. This means we can stop merging as soon as we reach the median.
2. Two Pointers, Track Only What We Need:
   We use two pointers to go through both arrays. We keep a count of how many elements we’ve seen so far. When we reach the positions of the median (middle elements), we save their values.
3. Median Calculation:
   • For an odd total number, the median is the middle element.
   • For an even total number, the median is the average of the two middle elements.

Why does this work?
You don’t need the whole merged array to find the median—just the elements at the median positions.

When should you use this?
This approach is more efficient in terms of space and slightly faster since it stops early.

Code Implementation

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        n1, n2 = len(nums1), len(nums2)
        n = n1 + n2
        ind2 = n // 2
        ind1 = ind2 - 1
        cnt = 0
        ind1el, ind2el = -1, -1

        i, j = 0, 0
        while i < n1 and j < n2:
            if nums1[i] < nums2[j]:
                if cnt == ind1:
                    ind1el = nums1[i]
                if cnt == ind2:
                    ind2el = nums1[i]
                cnt += 1
                i += 1
            else:
                if cnt == ind1:
                    ind1el = nums2[j]
                if cnt == ind2:
                    ind2el = nums2[j]
                cnt += 1
                j += 1

        while i < n1:
            if cnt == ind1:
                ind1el = nums1[i]
            if cnt == ind2:
                ind2el = nums1[i]
            cnt += 1
            i += 1

        while j < n2:
            if cnt == ind1:
                ind1el = nums2[j]
            if cnt == ind2:
                ind2el = nums2[j]
            cnt += 1
            j += 1

        if n % 2 == 1:
            return float(ind2el)
        return (ind1el + ind2el) / 2.0

Optimal Solution (Binary Search Partition)

Intuition and Approach

How does this solution work?

1. Key Idea – Partition the Arrays:
   The goal is to partition the two arrays into two halves such that:
   • The left half contains the same number of elements as the right half (or one more if the total is odd).
   • Every element in the left half is less than or equal to every element in the right half.
2. Why Partition?
   If you can split the arrays this way, the median will be either:
   • The maximum of the left half (if odd total length).
   • The average of the maximum of the left half and the minimum of the right half (if even total length).
3. How to Find the Partition?
   • Always binary search the smaller array for efficiency.
   • For a given position in the first array, you know exactly where the cut should be in the second array so that the left and right halves have the correct number of elements.
   • Use binary search to find the correct partition where:
     • The largest value on the left of the partition in the first array is less than or equal to the smallest value on the right of the partition in the second array.
     • And vice versa.
4. Why does this work?
   This method ensures you always find the correct median in O(log(min(n, m))) time, which is much faster than merging.

When should you use this?
This is the best approach for large arrays or when you want the most efficient solution.

Code Implementation

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        n1, n2 = len(nums1), len(nums2)
        # Always binary search on the smaller array
        if n1 > n2:
            return self.findMedianSortedArrays(nums2, nums1)
        
        n = n1 + n2
        left = (n1 + n2 + 1) // 2  # Number of elements in left half
        low, high = 0, n1
        
        while low <= high:
            mid1 = (low + high) // 2
            mid2 = left - mid1
            
            l1 = float('-inf') if mid1 == 0 else nums1[mid1 - 1]
            l2 = float('-inf') if mid2 == 0 else nums2[mid2 - 1]
            r1 = float('inf') if mid1 == n1 else nums1[mid1]
            r2 = float('inf') if mid2 == n2 else nums2[mid2]
            
            # Check if we have a valid partition
            if l1 <= r2 and l2 <= r1:
                if n % 2 == 1:
                    return float(max(l1, l2))
                else:
                    return (max(l1, l2) + min(r1, r2)) / 2.0
            elif l1 > r2:
                high = mid1 - 1
            else:
                low = mid1 + 1
        return 0.0  # Should not reach here

Final Thoughts

The “Median of Two Sorted Arrays” problem is a fantastic example of how you can optimize a solution step by step:

• Start with brute force to get the basics.
• Use a better approach to improve space and time.
• Master the optimal binary search partition for the fastest solution.

If you understand the intuition behind each method, you’ll be able to tackle similar problems with confidence.

For any changes to the article, kindly email at code@codeanddebug.in or contact us at +91-9712928220.