Median of Two Sorted Arrays

2022-03-21 19:30:38 +08:00
parent 29138c91a3
commit 35c715ebb1
3 changed files with 91 additions and 0 deletions
--- a/src/s0004_median_of_two_sorted_arrays.cpp
+++ 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<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;
+  }
+}