问题分析
对于寻找两个正序数组的中位数,可以进一步普适化成寻找两个正序数组中第 k 小的数。面对这类问题,可以使用分治的思路,将查找的时间复杂度降低到 O ( l o g ( m + n ) ) O(log(m+n)) O(log(m+n))
算法描述
分:在两个数组中找到第 k 小的数,可以分解成比较两个数组的第 k 2 \frac{k}{2} 2k 个元素
治:记两个数组nums1
和nums2
的起始下标为idx1
和idx2
,两个数组的第
k
2
\frac{k}{2}
2k 个元素的下标为t1
和t2
- 若
nums1[t1] <= nums2[t2]
,说明nums1
位于t1
左区间的元素不可能是第 k 小的数 - 若
nums1[t1] > nums2[t2]
,说明nums2
位于t1
左区间的元素不可能是第 k 小的数
合并:排除了一些不可能的元素后,重新更新 k k k 和数组的起始下标,缩小搜索范围
- 若
nums1[t1] <= nums2[t2]
,nums1
的起始下标从t1 + 1
开始,原问题变为找两个数组第k - (t1 - idx + 1)
个元素 - 若
nums1[t1] > nums2[t2]
,nums2
的起始下标从t2 + 1
开始,原问题变为找两个数组第k - (t2 - idx + 1)
个元素
程序代码
C++
class Solution {
private:
double findKth(vector<int>& nums1, vector<int>& nums2, int k) {
int m = nums1.size(), n = nums2.size();
int idx1 = 0, idx2 = 0;
while( true ) {
if( idx1 == m ) return nums2[idx2 + k - 1];
if( idx2 == n ) return nums1[idx1 + k - 1];
if( k == 1 ) return min(nums1[idx1], nums2[idx2]);
int t1 = min(idx1 + k / 2 - 1, m - 1);
int t2 = min(idx2 + k / 2 - 1, n - 1);
if( nums1[t1] <= nums2[t2] ) {
k -= t1 - idx1 + 1;
idx1 = t1 + 1;
}
else {
k -= t2 - idx2 + 1;
idx2 = t2 + 1;
}
}
}
public:
double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
int m = nums1.size(), n = nums2.size();
if( (m + n) & 1 ) {
return findKth(nums1, nums2, (m + n + 1) / 2);
}
double t1 = findKth(nums1, nums2, (m + n) / 2);
double t2 = findKth(nums1, nums2, (m + n) / 2 + 1);
return (t1 + t2) / 2;
}
};
Go
func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
m, n := len(nums1), len(nums2)
if (m + n) % 2 == 1 {
return float64(findKth(nums1, nums2, (m + n + 1) / 2))
}
t1 := float64(findKth(nums1, nums2, (m + n) / 2))
t2 := float64(findKth(nums1, nums2, (m + n) / 2 + 1))
return (t1 + t2) / 2
}
func findKth(nums1 []int, nums2 []int, k int) int {
m, n := len(nums1), len(nums2)
idx1, idx2 := 0, 0
for {
if idx1 == m {
return nums2[idx2 + k - 1]
}
if idx2 == n {
return nums1[idx1 + k - 1]
}
if k == 1 {
return min(nums1[idx1], nums2[idx2])
}
t1 := min(idx1 + k/2 - 1, m - 1);
t2 := min(idx2 + k/2 - 1, n - 1);
if nums1[t1] <= nums2[t2] {
k -= t1 - idx1 + 1
idx1 = t1 + 1
} else {
k -= t2 -idx2 + 1
idx2 = t2 + 1
}
}
}
func min(x, y int) int {
if x <= y {
return x
}
return y
}