Find the Median of Two Sorted Arrays - Leet Code Solution

Problem Statement

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

1. Simple Solution

Lets think of a simple approach. We need the numbers ordered in the sequence, and we can easily fetch the median number. Since the arrays are sorted. We can merge the two sorted array by using a simple merge algorithm.

Then, we just have to get the middle index, if the length of merged array is odd. Else, we need the average of two middle indexes. lets look at the code:

public class SimpleSolution {
    private int[] arr1;
    private int[] arr2;
    
    public SimpleSolution(int[] arr1, int[] arr2) {
        this.arr1 = arr1;
        this.arr2 = arr2;
    }

    private int[] merge() {
        int[] res = new int[this.arr1.length + this.arr2.length];
        //index for arr1
        int i=0;
        //index for arr2
        int j=0;
        // result index
        int k=0;

        while (i < this.arr1.length && j < this.arr2.length) {
            if (this.arr1[i] < this.arr2[j]) {
                res[k] = this.arr1[i];
                i++;
            }
            else {
                res[k] = this.arr2[j];
                j++;
            }
            k++;
        }
        for (; i<this.arr1.length; i++) {
            res[k] = this.arr1[i];
            k++;
        }
        for (; j<this.arr2.length; j++) {
            res[k] = this.arr2[j];
            k++;
        }

        return res;
    }

    public double getMedian() {
        int[] mergedArray = this.merge();
        int l = mergedArray.length;
        
        //if length is odd. Median is middle number
        //else, its average of middle two elements

        if (l % 2 != 0) {
            return mergedArray[l/2];
        }

        //else
        return (double)(mergedArray[(l-1)/2] + mergedArray[(l)/2]) / 2;
    }
}

Complexity

We are taking an extra space equals to length of two arrays. Space complexity: O(m + n). m and n are length of two arrays. Runtime complexity is around O(m + n).

2. Optimized Solution

Lets look at the definition of median once again. We want an element (or two elements if count is even) with which every element on left side is less than. And, is greater than every element on right side.

The objective is to find an index in both the arrays which creates a partition in both arrays. Such that the elements on the partition index is the median.

Partitioning of arrays

The idea is to have a partition from both arrays, a kind of window. Which forms a left partition and right partition from both arrays. Where every element from left partition is less than every element from right partition.

A(A1, A2, A3, A4, A5) is 1st array
B(B1, B2, B3, B4, B5, B6) is 2nd array

After correct partition:
A1 A2 A3    A4 A5
B1 B2 B3    B4 B5 B6

Such that, A3 is less than equal to A4 or B4
And, B3 is also less than equal to A4 or B4

Since the array is sorted, we can say
Max(A3, B3) <= Min(A4, B4)

We can say above partition as correct partition only when both partition have equal size (left side may have an extra element), and which satisfies above condition.

Binary Search

We can partition our shorter array into half, just like binary search. And, we have to partition other array according to total elements and 1st partition.

Calculations

First we calculate half number of elements from total of both arrays.

l1 - size of of 1st array
l2 - size of 2nd array

halfElements = (l1 + l2 + 1)/2;
partitionA = (0, l1)/2
partitionB = haldElements - partitionA

Algorithm

We just want to iterate over shorter array, and partition over 2nd array will vary according to 1st array’s partition. Since we are maintaining equal number of elements in both partitions. We just need to satisfy the condition where

Max(left partition) <= Min(Right partition)

And, there are two other conditions as well. If last elements from Array-1 from left partition is greater than 1st element of Array-2 in right partition, we need to move left in Array-1.

Similarly, if last element from Array-2 from left partition is greater than 1st element of Array-1 in right partition, we need to move right in Array-1.

And just keep on moving like this, calculating partition index again and again. We will get the result.

Correct partition

When we found the correct partition. We need to check if the total (l1 + l2) is odd, the result is max element in left partition. Else, the result is average of max of left partition, and min of right partition.

If (l1 + l2) is odd
Result: Max (left partition)

If (l1 + l2) is even
Result: Average(max(left-partition), min(right-partition))

Code

public class OptimizedSolution {
    private int[] arr1;
    private int[] arr2;

    public OptimizedSolution(int[] arr1, int[] arr2) {
        this.arr1 = arr1;
        this.arr2 = arr2;
    }

    public double getMedian() {
        //lets take shorter array first
        if (this.arr1.length > this.arr2.length) {
            //swap references
            int[] t = this.arr1;
            this.arr1 = this.arr2;
            this.arr2 = t;
        }

        int l1 = this.arr1.length;
        int l2 = this.arr2.length;

        //need to partition shorter array first, then 2nd so that both partition have equal elements (+- 1)
        int ps = 0;     //partition start index
        int pe = l1;    //partition end index (Note it is not the last index)

        int halfElements = (l1 + l2 + 1)/2;

        while (ps <= pe) {
            //index of partition of 1st array
            int partitionA = (ps + pe) / 2;

            //calculate partition index of 2nd array
            int partitionB = halfElements - partitionA;

            //compare elements
            if (partitionA > 0 && this.arr1[partitionA-1] > this.arr2[partitionB]) {
                pe = partitionA - 1;
            }
            else if (partitionA < l1 && this.arr2[partitionB-1] > this.arr1[partitionA]) {
                ps = partitionA + 1;
            } 
            else {
                //found balanced both partitions
                int maxLeft = this.getMaxOfLeftPartition(partitionA, partitionB);
                //get median
                if ((l1 + l2) % 2 != 0) {
                    return maxLeft;
                }

                int minRight = this.getMinOfRightPartition(partitionA, partitionB);
                return (maxLeft + minRight)/2;
            }
        }

        return 0.0;
    }

    private int getMaxOfLeftPartition(int partitionA, int partitionB) {
        if (partitionA == 0) return this.arr2[partitionB-1];
        if (partitionB == 0) return this.arr1[partitionA-1];
        return Math.max(this.arr1[partitionA-1], this.arr2[partitionB-1]);
    }

    private int getMinOfRightPartition(int partitionA, int partitionB) {
        if (partitionA == this.arr1.length) return this.arr2[partitionB];
        if (partitionB == this.arr2.length) return this.arr1[partitionA];
        return Math.min(this.arr1[partitionA], this.arr2[partitionB]);
    }
}