Median of Two Sorted Arrays

Problem

You are given two integer arrays nums1 and nums2 of size m and n respectively, where each is sorted in ascending order. Return the median value among all elements of the two arrays.

Your solution must run in O(log(m+n)) time.

Examples

Example 1:

Input: nums1 = [1,2], nums2 = [3]

Output: 2.0

Explanation: Among [1, 2, 3] the median is 2.

Example 2:

Input: nums1 = [1,3], nums2 = [2,4]

Output: 2.5

Explanation: Among [1, 2, 3, 4] the median is (2 + 3) / 2 = 2.5.

Constraints

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

You should aim for a solution with O(log(min(n, m))) time and O(1) space, where n is the size of nums1 and m is the size of nums2.

Solution

Idea is to use the fact that we know how many elements would be on the left side of the merged array. We perform binary search on the smaller of the two arrays in terms of how many elements to select from it. If we take all the elements from A, then we know the solution lies left_size - M within B, which is given by j.

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        constexpr auto INF  = std::numeric_limits<int>::max();
        constexpr auto NINF = std::numeric_limits<int>::lowest();
        // Binary search on the smaller of the two vectors
        if (nums1.size() > nums2.size()) {
            return findMedianSortedArrays(nums2, nums1);
        }
 
        const std::vector<int> &A = nums1;
        const std::vector<int> &B = nums2;
        int M = A.size();
        int N = B.size();
 
        // We want to simulate partitioning the merged array
        // The left partition will have half the elements (+1 if odd length)
        int left_size = (M + N + 1) / 2;
 
        // Binary search is over how many elements we take from A
        // i = number of elements from A placed on the left side
        // j = number of elements from B placed on the left side
        int low = 0;
        int high = M;
        while (low <= high) {
            // Left side contains left_size elements, therefore j = left_size - i
            int i = (low + high) / 2;
            int j = left_size - i;
            
            // Partitions:
            // A: [A[0], ..., A[i-1] | A[i], ..., A[M-1]]
            // B: [B[0], ..., B[j-1] | B[j], ..., B[N-1]]
            // If i = 0, A contributes nothing to the left side
            // If i = M, A contributes everything to the left side
            int leftA = (i == 0) ? NINF : A[i - 1];
            int rightA = (i == M) ? INF : A[i];
            int leftB = (j == 0) ? NINF : B[j - 1];
            int rightB = (j == N) ? INF : B[j];
 
            // We found the correct partition if everything on left is 
            //   <= everything on the right
            // Since A/B are each sorted, we only need to check the crossing boundaries
            if (leftA <= rightB && leftB <= rightA) {
                auto leftMax = std::max(leftA, leftB);
                auto rightMin = std::min(rightA, rightB);
                // Odd length, since left has one more element,
                // median is largest element on the left
                if ((M + N) % 2 == 1) { return static_cast<double>(leftMax); }
                // Even length, median is average of largest element on left
                //   and smallest element on right
                return static_cast<double>(leftMax + rightMin) / 2.0;
            }
            
            // If leftA > rightB, we took too many elements from A, so move i left
            // Otherwise, we took too few elements from A, so move i right
            if (leftA > rightB) { high = i - 1;}
            else { low = i + 1; }
        }
        return 0;
    }
};