LeetCode 2 – Median of Two Sorted Arrays

Posted on December 21, 2013

Last updated on December 21, 2013

Solution to LeetCode Median of Two Sorted Arrays problem.

The solution and its proof are non-trivial, so I would refer you to “Pearls of Functional Algorithm Design”, chapter “A Selection Problem” for an in-depth explanation of computing this. The algorithm there gives us a general way to compute the \(k\)’th smallest number given two sorted arrays. The algorithm is quite neat, although this array version unfortunately suffers from indexitis – there are a lot of indices floating around.

public double findMedianSortedArrays(int A[], int B[]) {
  int total_length = A.length + B.length;
  int median_index = total_length/2;
  double median = (double) kthSmallest(A,B,median_index, 0, A.length, 0, B.length);

  if(total_length % 2 == 0 && total_length >= 2){
    median = (median + (double) kthSmallest(A,B,median_index-1, 0, A.length, 0, B.length))/2;
  }
  return median;
}

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);
    }
  }
}
Median of Two Sorted Arrays
Markdown SHA1: 6902a9f101f880060b722e8212f6b31197530bbc