LeetCode: Median of Two Sorted Arrays: explanation and solution with C++/Java/Python

Problem

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

There are two sorted arrays nums1 and nums2 of size m and n respectively.

Find the median of the two sorted arrays.

The overall run time complexity should be $O(log(m+n))$.

You may assume nums1 and nums2 cannot be both empty.

Example 1:

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

The median is 2.0

Example 2:

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

The median is (2 + 3)/2 = 2.5

Explanation

To find the median of an array, there are two sub-cases here according to the size of array(assume index start from 0):

When length is odd:
   median = arr[size/2]

When length is even:
   median = (arr[(size-1)/2] + arr[size/2]) / 2

It’s will be comre complicated to compute the median of two array.

Naive solution

We can merge the two array into one array, keep the sorted property of array. The time complexity for this will be $O(max(m, n))$, with a constant time for computing the median.

 class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        vector<int> res;
        int i=0;
        int j=0;
        int mid = (nums1.size() + nums2.size()) / 2;
        int is_odd = ((nums1.size() + nums2.size()) % 2);
        int end = is_odd ? mid : mid+1;
        while(i < nums1.size() && j < nums2.size() && res.size() <= end) { 
            int t;
            if(nums1[i] > nums2[j]) {
                t = nums2[j];
                j++;
            } else { 
                t = nums1[i];
                i++;
            }
            res.push_back(t);
        }
        while( i < nums1.size() && res.size() <= end) {
            res.push_back(nums1[i]);
            i++;
        }
        while( j < nums2.size() && res.size() <= end) {
            res.push_back(nums2[j]);
            j++;
        }

        if(is_odd) return (double)res[mid];
        else return ((double)res[mid] + double(res[mid-1])) / 2.0;
    }
};

Optimization to log(m+n)

For almost any algorithms which perform on array and complexity is $O(logN)$, we need to split the size into almost size/2 during every iteration. Think about the binary search for a sorted array, it’s complexity is $O(logN)$.

If we could find an algorithm which find the Kth largest element in two sorted array, then use a split strategy we could come out this pseudocode:

find_median_sorted_array(A, B) {
  len = size(A) + size(B)
  if (len % 2) == 1 {
     find_kth(A, B, len/2 + 1)
  } else { 
     ma = find_kth(A, B, len/2)
     mb = find_kth(A, B, len/2 + 1)
     return (ma + mb) / 2
  }
}

Find Kth largest in one sorted array

Algorithm design guide:
Try to solve a similar but much easy case, then expand it into harder cases.

Let’s first solve this much easier and similar problem. Like the Binary Search algorithm, we first check the compare result of size/2 with K, if (size/2) >= K, we can drop the right part of array. Otherwise we can drop the left part and find K-(size/2+1) th in right part of this array.

The fundamental idea is reducing the size to half in every iteration.

Find Kth largest in two sorted array, without merging

This problem is similar with Kth largest in one array, but non-trivial.

There are 4 sub-cases which depends on the compare result of A[m/2] and B[n/2]. Let’s assume i = m/2 and j = n/2, A[i] <= B[j].

      left_part          |        right_part
A[0], A[1], ..., A[i-1]  |  A[i], A[i+1], ..., A[m-1]
B[0], B[1], ..., B[j-1]  |  B[j], B[j+1], ..., B[n-1]

If (i+j) >= K, this means there are more than K elements in left part of two array, so we can drop the right part of B. Otherwise, if (i+j) < K, we can drop the right part of A and try to find the K-(m/2+1) th largest in left parts.

In this algorithm we almost reduce 1/4 of the total size of two array. The complexity analysis is very hard for most programmers. Chapter “A Selection Problem” of Pearls of Functional Algorithm Design contains the detail description.

C/C++

This is C version solution, take the advantage of pointer, the parameter of “int A[]” with index operation “A+index” point to sub-array, this save two parameters in findMedianSortedArrays. For C++ version, we need to add a start index for array, similar with below Java solution.

Be careful with those index in edge cases!

  #define min(a, b) ((a) > (b) ? (b) : (a))

  //find the kth from A and B
  int findKth(int A[], int m, int B[], int n, int k) {
    if(m <= 0) return B[k];
    if(n <= 0) return A[k];
    if(B[n/2] >= A[m/2]) {
      if((m/2 + n/2) >= k)
        return findKth(A, m, B, n/2, k);
      else
        return findKth(A + m/2 + 1, m - (m/2 + 1), B, n, k - (m/2 + 1));
    } else {
      if((m/2 + n/2) >= k)
        return findKth(A, m/2, B, n, k);
      else
        return findKth(A, m, B + (n/2 + 1), n - (n/2 + 1), k - (n/2 + 1));
    }
  }

  double findMedianSortedArrays(int A[], int m, int B[], int n) {
    int len = m + n;
    if((len % 2) == 1) {
      return findKth(A, m, B, n, len/2);
    } else {
      double a = (double)findKth(A, m, B, n, (len-1)/2);
      double b = (double)findKth(A, m, B, n, len/2);
      return (a + b)/2.0;
    }
}

Java

class Solution {
    public double findMedianSortedArrays(int A[], int B[]) {
        int len = A.length + B.length;
        int index = len/2;

        if(len % 2 == 1) {
            return (double) kthSmallest(A, B, index, 0, A.length, 0, B.length);
        } else {
            return ((double) kthSmallest(A, B, index, 0, A.length, 0, B.length) +
                    (double) kthSmallest(A, B, index-1, 0, A.length, 0, B.length)) / 2;
        }
    }

    private int kthSmallest(int[] A, int[] B, int k, int aL, int aR, int bL, int bR){
        if(aL == aR) return B[bL+k];
        if(bL == bR) return A[aL+k];
        else{
            int midA = (aL+aR)/2;
            int midB = (bL+bR)/2;
            int lenA = midA - aL;
            int lenB = midB - bL;

            if(A[midA] <= B[midB]){
                if(k <= lenA + lenB)
                    return kthSmallest(A, B, k, aL, aR, bL, midB);
                else
                    return kthSmallest(A, B, k-lenA-1, midA+1, aR, bL, bR);
            }
            else{
                if(k <= lenA + lenB)
                    return kthSmallest(A, B, k, aL, midA, bL, bR);
                else
                    return kthSmallest(A, B, k-lenB-1, aL, aR, midB+1, bR);
            }
        }
    }
}

Preparing for an interview? Checkout this!