Find the Median of Two Sorted Arrays (via Leetcode)¶
Date published: 2020-03-08
Category: Python
Subcategory: Advanced Algorithms
Tags: functions, loops, sets
This example was found on Leetcode.
In this tutorial, I'll walk extensively through my logic for solving this problem.
I'll solve it the brute force way and later solve it in a more performant way.
Problem statement:
There are two sorted arrays, nums1 and nums2 of size m and n respectively.
Find the median of the two sorted arrays. Assume both arrays are non-empty.
The overall run time complexity should be O(log (m+n)).
Test cases:
# odd-numbered lists; median is middle number
# assert find_median_of_two_sorted_arrays([1], [1, 2]) == 1
# assert find_median_of_two_sorted_arrays([1], [1, 2, 3, 4]) == 2
# assert find_median_of_two_sorted_arrays([1, 5], [0, 3, 6]) == 3
# even-numbered lists; median is average of two "middle" numbers
# assert find_median_of_two_sorted_arrays([1], [2]) == 1.5
# assert find_median_of_two_sorted_arrays([1], [2, 4, 6]) == 3
# assert find_median_of_two_sorted_arrays([1, 3], [2, 4]) == 2.5
Brute Force Method¶
Pseudocode¶
I want to calculate the "middle" number(s) used to calculate the median. To find the "middle" number(s), I'll use the "merge step" pattern - without creating a merged list. This pattern use two pointers - one for each list - to track the index at each list. The elements at each index are compared, and the pointer is incremented for the smaller value.
The median for an even-numbered list such as [1, 2, 3, 4] is the average of the two "middle" numbers of 2 and 3 which is 2.5.
The median for an odd-number list such as [1, 2, 3] is the "middle" number of 2.
That's enough to begin coding.
Coded Solution¶
def find_median_of_two_sorted_arrays(nums1, nums2):
nums1_length, nums2_length = len(nums1), len(nums2)
total_length = nums1_length + nums2_length
median_index2 = total_length // 2
nums1_index, nums2_index = 0, 0
previous_value, current_value = None, None
joined_list_index = 0
while joined_list_index <= median_index2:
# helps us store the "middle" numbers to use in median final calculations
previous_value = current_value
# checks if there are still elements left in both arrays. If so, we need to compare the elements at the
# current positions of the pointers and move the pointer with the smaller value.
if nums1_index < nums1_length and nums2_index < nums2_length:
if nums1[nums1_index] < nums2[nums2_index]:
current_value = nums1[nums1_index]
nums1_index += 1
else:
current_value = nums2[nums2_index]
nums2_index += 1
# condition checks if elements still left in first array
elif nums1_index < nums1_length:
current_value = nums1[nums1_index]
nums1_index += 1
# no elements left in first array
else:
current_value = nums2[nums2_index]
nums2_index += 1
joined_list_index += 1
# do median calculation
if total_length % 2 == 0:
return (previous_value + current_value) / 2.0
else:
return current_value
All the test cases below pass.
assert find_median_of_two_sorted_arrays([1], [1, 2]) == 1
assert find_median_of_two_sorted_arrays([1], [1, 2, 3, 4]) == 2
assert find_median_of_two_sorted_arrays([1, 5], [0, 3, 6]) == 3
assert find_median_of_two_sorted_arrays([1], [2]) == 1.5
assert find_median_of_two_sorted_arrays([1], [2, 4, 6]) == 3
assert find_median_of_two_sorted_arrays([1, 3], [2, 4]) == 2.5