47 lines
1.3 KiB
C++
47 lines
1.3 KiB
C++
#include "s0004_median_of_two_sorted_arrays.hpp"
|
|
|
|
int S0004::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 S0004::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;
|
|
}
|
|
}
|