diff --git a/include/s0004_median_of_two_sorted_arrays.hpp b/include/s0004_median_of_two_sorted_arrays.hpp new file mode 100644 index 0000000..5305e1b --- /dev/null +++ b/include/s0004_median_of_two_sorted_arrays.hpp @@ -0,0 +1,26 @@ +#ifndef S0004_MEDIAN_OF_TWO_SORTED_ARRAYS +#define S0004_MEDIAN_OF_TWO_SORTED_ARRAYS + +#include + +class Solution { + public: + /** + * @brief Median of Two Sorted Arrays + * + * Given two sorted arrays `nums1` and `nums2` of size `m` and `n` + * respectively, return the median of the two sorted arrays. + * + * The overall run time complexity should be `O(log(m+n))`. + * + * @param nums1 the first sorted array + * @param nums2 the second sorted array + * @return the median of the two sorted arrays + */ + double findMedianSortedArrays(std::vector& nums1, + std::vector& nums2); + int getKthElement(const std::vector& nums1, + const std::vector& nums2, int k); +}; + +#endif diff --git a/src/s0004_median_of_two_sorted_arrays.cpp b/src/s0004_median_of_two_sorted_arrays.cpp new file mode 100644 index 0000000..06fc867 --- /dev/null +++ b/src/s0004_median_of_two_sorted_arrays.cpp @@ -0,0 +1,46 @@ +#include "s0004_median_of_two_sorted_arrays.hpp" + +int Solution::getKthElement(const std::vector& nums1, + const std::vector& nums2, int k) { + int m = nums1.size(); + int n = nums2.size(); + int index1 = 0, index2 = 0; + + while (true) { + // Edge case + if (index1 == m) { + return nums2[index2 + k - 1]; + } + if (index2 == n) { + return nums1[index1 + k - 1]; + } + if (k == 1) { + return std::min(nums1[index1], nums2[index2]); + } + + // Normal case + int newIndex1 = std::min(index1 + k / 2 - 1, m - 1); + int newIndex2 = std::min(index2 + k / 2 - 1, n - 1); + int pivot1 = nums1[newIndex1]; + int pivot2 = nums2[newIndex2]; + if (pivot1 <= pivot2) { + k -= newIndex1 - index1 + 1; + index1 = newIndex1 + 1; + } else { + k -= newIndex2 - index2 + 1; + index2 = newIndex2 + 1; + } + } +} + +double Solution::findMedianSortedArrays(std::vector& nums1, + std::vector& nums2) { + int totalLength = nums1.size() + nums2.size(); + if (totalLength % 2 == 1) { + return getKthElement(nums1, nums2, (totalLength + 1) / 2); + } else { + return (getKthElement(nums1, nums2, totalLength / 2) + + getKthElement(nums1, nums2, totalLength / 2 + 1)) / + 2.0; + } +} diff --git a/tests/s0004_median_of_two_sorted_arrays.cpp b/tests/s0004_median_of_two_sorted_arrays.cpp new file mode 100644 index 0000000..d1f509e --- /dev/null +++ b/tests/s0004_median_of_two_sorted_arrays.cpp @@ -0,0 +1,19 @@ +#include "s0004_median_of_two_sorted_arrays.hpp" + +#include + +TEST(Problem4, Case1) { + std::vector nums1{1, 3}; + std::vector nums2{2}; + double answer{2}; + Solution solution; + EXPECT_EQ(solution.findMedianSortedArrays(nums1, nums2), answer); +} + +TEST(Problem4, Case2) { + std::vector nums1{1, 2}; + std::vector nums2{3, 4}; + double answer{2.5}; + Solution solution; + EXPECT_EQ(solution.findMedianSortedArrays(nums1, nums2), answer); +}