LeetCode — Median of Two Sorted Arrays

Problem definition is here.

Given two sorted arrays nums1 and nums2 of size m and n respectively, return the median of the two sorted arrays.

Sample inputs

Input: nums1 = [1,3], nums2 = [2]
Output: 2.00000
Explanation: merged array = [1,2,3] and median is 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.

TOC

• What is median?
• Draft
• Hypothesis
  • Iteration
  • Picking elements
  • Handling odd/even number of elements
• Solution

What is median?

In a given sorted data set, median is the middle element. It is not neccessarily the mean/average.

[2, 2, 3, 4, 10]

⇨ mean = (2+2+3+4+10)/5 = 4.2

⇨ median 3

...

Draft

Few things to note, before solving the problem:

1. Input arrays are already sorted. That means, we don't have to sort the resulting array if we pick the elements in correct order.
2. Result asking for median of the resulting array, not the array itself. So, we don't have to allocate a new arrray at all.
3. From the inputs, if total number of elements is odd, we'll pick a single element. Otherwise, average of the two elements in the middle.

...

Hypothesis

Iteration

Since we're looking for median, we don't have to iterate all the way upto m + n. Here, m and n are size of each array. For odd number of elements, we pick one median and it's two for even numbers. So, how far do we iterate?

Notice that I mentioned median index and index starts with zero. So, when I mention first index, it is the second element in the result array.

Can we formulate the number here?
4 ⇨ 4/2 = 2 [m1 = 1, m2 = 2]
5 ⇨ 5/2 = 2 [m2 = 2]

For a given total(m+n), we'll iterate up to (m+n)/2 index. If the total is odd, pick only m2, otherwise average of m1 + m2. So, last two iterations matters most.

Picking elements

In order to fill elements in result array, we have to pick the least element of the two arrays.

Input:
[1, 4]
[2, 9]

0 --> [1]
1 --> [1,2]
2 --> [1,2,4]
3 --> [1,2,4,9] // we don't move this far — iteration stops at index 2

What if same number present in two arrays? 🤔 It doesn't matter. Both will be placed adjacent in array anyway, and we're just focusing on the value.

So, we need to keep two pointers(i, j) for both arrays to pick element in head. What do we pick?

if (nums1[i] < nums2[j]) {
   m2 = nums2[i]
   i++
} else {
   m2 = nums2[j]
   j++
}

There is a special case as well. What if one of the arrays ran out of elements?. In such case, we have to pick element from the other array and move pointer accordingly.

// num1 reached its end
if (num1.size == i) {
   m2 = nums2[j]
   j++
}

Handling odd/even number of elements

Through each iteration, we update medians m1 & m2. At the end, last two picks (or just one) considered for median. Values always stored in m2 and on each iteration, m1 will cache the previous value of m2. Things will be much clear with actual solution below.

...

Solution

Putting all the above together, we have the solution here. Again, I added comments inline.

fun findMedianSortedArrays(nums1: IntArray, nums2: IntArray): Double {

    // Pointers to array
    var i = 0
    var j = 0

    // Identified medians
    var m1 = 0
    var m2 = 0

    // iterate upto - (m+n)/2
    for (k in 0..(nums1.size+nums2.size)/2) {
        // caching previous value
        m1 = m2
        if (i == nums1.size) {
            // First array ran out of elements, pick second
            m2 = nums2[j]
            j++
        } else if (j == nums2.size) {
            // Second array ran out of elements, pick first
            m2 = nums1[i]
            i++
        } else if (nums1[i] < nums2[j]) {
            // First array element is lower
            m2 = nums1[i]
            i++
        } else {
            // Second array element is lower
            m2 = nums2[j]
            j++
        }
    }

    // For even numbers avg of m1 & m2, m2 otherwise
    return if ((nums1.size+nums2.size)%2 == 0) {
        (m1 + m2) / 2.0
    } else m2.toDouble()    
}

🧑‍💻🧑‍💻 peace 🧑‍💻🧑‍💻