Median of Two Sorted Arrays

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 <vector>
+
+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<int>& nums1,
+                                std::vector<int>& nums2);
+  int getKthElement(const std::vector<int>& nums1,
+                    const std::vector<int>& nums2, int k);
+};
+
+#endif

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<int>& nums1,
+                            const std::vector<int>& 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<int>& nums1,
+                                        std::vector<int>& 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;
+  }
+}

tests/s0004_median_of_two_sorted_arrays.cpp

@@ -0,0 +1,19 @@
+#include "s0004_median_of_two_sorted_arrays.hpp"
+
+#include <gtest/gtest.h>
+
+TEST(Problem4, Case1) {
+  std::vector<int> nums1{1, 3};
+  std::vector<int> nums2{2};
+  double answer{2};
+  Solution solution;
+  EXPECT_EQ(solution.findMedianSortedArrays(nums1, nums2), answer);
+}
+
+TEST(Problem4, Case2) {
+  std::vector<int> nums1{1, 2};
+  std::vector<int> nums2{3, 4};
+  double answer{2.5};
+  Solution solution;
+  EXPECT_EQ(solution.findMedianSortedArrays(nums1, nums2), answer);
+}