LeetCode 4 - Median of Two Sorted Arrays

作者 QIFAN 日期 2017-03-05
LeetCode 4 - Median of Two Sorted Arrays

原题链接: 4. Median of Two Sorted Arrays


题干

给定两个已排序的数组 nums1 和 nums2 ,长度分别为 m 和 n ,返回这 m + n 个数的中位数。要求时间复杂度为 $O(log(m + n))$ 。

思路

一开始一直卡在 binary search 想 base case ,看了 discussion 以后的启发,就是 findKth 的一个延伸。如果 m + n 是奇数,则 k = (m + n + 1) / 2 ,如果 m + n 是偶数,则 k = (m + n + 1) / 2 与 (m + n + 2) / 2

两个数组都从起始位置也就是 0 开始。(nums1, 0, nums2, 0, k)
如果 nums1 的长度大于 nums2 ,(nums2, 0, num1, 0, k)

如果 k == 1 返回 min(nums1[0], nums2[0])

i = min(0 + k / 2, nums1.length)
j = k / 2
如果 nums1[i] < nums2[j] ,那 nums1 中的 0 - i 一定在前 k 个里,(num1, i, nums2, 0, k - i)
如果 nums1[i] > nums2[j] ,那 nums2 中的 0 - j 一定在前 k 个里,(num1, 0, nums2, j, k - j)

概括一些,假设分别从 nums1 的第 i1 个和 nums2 的第 i2 个开始找第 k 个元素 (nums1, i1, nums2, i2, k):
如果 i1 == nums1.length ,返回 nums2[i2 + k]
如果 i2 == nums2.length ,返回 nums1[i1 + k]
if (nums2.length - i2 < nums1.length - i1) , (nums2, i2, nums1, i1, k)

java 不像 c 可以直接用指针表示数组,所以从后往前写比较方便,就是找到第 k 个大的数字。基本逻辑不变,只是 k = 1 时返回的是更大的数,缩小范围时把较大的一部分截掉。

代码:

public double findMedianSortedArrays(int[] nums1, int[] nums2) {
int totalL = nums1.length + nums2.length;
if (totalL % 2 == 1) {
return findKth(nums1, nums1.length, nums2, nums2.length, (totalL + 1) >> 1);
}
return (findKth(nums1, nums1.length, nums2, nums2.length, (totalL + 1) >> 1) + findKth(nums1, nums1.length, nums2, nums2.length, (totalL + 2) >> 1)) / 2.0;
}
// m 表示当前 nums1 的位置范围是 0-m
private int findKth(int[] nums1, int m, int[] nums2, int n, int k) {
if (m == 0) {
return nums2[n - k];
}
if (n == 0) {
return nums1[m - k];
}
if (k == 1) {
return Math.max(nums1[m - 1], nums2[n - 1]);
}
if (m > n) {
return findKth(nums2, n, nums1, m, k);
}
int i = Math.min(k / 2, m);
int j = k / 2;
if (nums1[m - i] > nums2[n - j]) {
return findKth(nums1, m - i, nums2, n, k - i);
}
return findKth(nums1, m, nums2, n - j, k - j);
}

优化的地方可以是,偶数中位数不用重新来过算后面一个,如果知道 (totalL + 1) / 2 时两个数组的范围,就返回上一个中的更大值就好了。