题目描述
给定两个大小为 m
和 n
的有序数组 nums1
和 nums2
。请你找出这两个有序数组的中位数,并且要求算法的时间复杂度为 O(log(m + n))
。
你可以假设 nums1
和 nums2
不会同时为空。
示例 1:
nums1 = [1, 3]
nums2 = [2]
则中位数是 2.0
示例2:
nums1 = [1, 2]
nums2 = [3, 4]
则中位数是 (2 + 3)/2 = 2.5
官方难度
Hard
解决思路
最简单粗暴的解决思路,从左到右遍历两个数组,找到中位数的位置:
- 计算
nums1
和nums2
的中位数,mid=(len(nums1)+len(nums2))/2
,如果为长度和为偶数,则为mid=(len(nums1)+len(nums2))/2
和mid-1
- 同时从左到右遍历两个数组,如果
nums1>nums2
,nums2
的下标++,否则nums1
的下标++,直到找到第mid个数。
根据上面的分析,我们可以很容易的得到直接方案流程如下:
- 初始化
start1
和start2
代表当前nums1
和nums2
的下标; - 遍历
mid+1
次,每次找第i
大的数,i
从0开始,最终要找到第mid
或者mid+1
的数。 - 注意
nums1
或者nums2
的下标越界问题。 - 如果
nums1[start1]>nums2[start2]
说明第i
大的数字为nums2[start2]
- 如果
i>=mid
,可以结束循环,根据两个数组长度总和计算返回的中位数。
基于这个流程,我们可以实现类似下面的代码:
class Solution(object):
def findMedianSortedArrays(self, nums1, nums2):
"""
:type nums1: List[int]
:type nums2: List[int]
:rtype: float
"""
n_len = len(nums1) + len(nums2)
mid = n_len / 2
start1 = 0
start2 = 0
pre = 0
current_val = 0
#当前第[0,i]大的数,最终要求第mid和mid+1位置的数字
for i in range(mid+1):
pre = current_val
if start1 > len(nums1)-1:
current_val = nums2[start2]
start2 += 1
elif start2 > len(nums2)-1:
current_val = nums1[start1]
start1 += 1
elif nums1[start1] < nums2[start2]:
current_val = nums1[start1]
start1 += 1
else:
current_val = nums2[start2]
start2 += 1
if i >= mid:
if n_len % 2 != 0:
return current_val
return float(current_val+pre)/2
优化
看到有序数组,马上想到二分查找的思路,假设中位数是求第mid
个数,我们假设在第一个数组中求mid/2
大的数字,第二个数组中求mid/2
大的数字,然后比较这两个数字,如果数组1的结果大,其实可以保证数组2的第mid/2
个数肯定是要小于两个数组合并后的第mid
个数,因此可以排除第二个数组的前mid/2
个数字,以输入nums1:[1,2,5,6], nums2:[3,5,6,7,9]
为例子,
1.要求mid=(4+5+1)/2=5
个数字,
- 从
nums1
中求第5/2大的数字,从nums2
中求第5/2大的数字, nums1[2-1] = 2, nums2[2-1] = 5
- 可以将nums1中的
1,2
两个数字都去掉,也就变成求整个数组的[5,6], [3,5,6,7,9]
第3大的数字。 - 一直求到第1大的数字时候,比较
nums1[0]和nums2[0]
的最小值则是最终结果。 - 需要注意数组越界,数组为空,下标和mid对应以及偶数长度数组和奇数长度数组的问题。
基于以上思路我们可以梳理流程如下:
- 计算
n
为nums1
的长度,m
为nums2
的长度。 - 解决偶数长度和奇数长度的问题,通过求
(n+m+1)/2和(n+m+2)/2
两个位置的数字,将这两类问题合并为一类问题,最终结果为[(n+m+1)/2 位置的数字+(n+m+2)/2位置的数字]*0.5
- 定义递归函数,用于求第K大的数字
- 递归函数定义结束条件,k==1
- 判断是否数组越界
- 比较大小,然后递归求
k - 去除数组的数量
基于这个流程,我们可以实现类似下面的代码:
class Solution(object):
def findMedianSortedArrays(self, nums1, nums2):
"""
:type nums1: List[int]
:type nums2: List[int]
:rtype: float
"""
n = len(nums1)
m = len(nums2)
if not nums1:
return (nums2[(n+m+1)/2-1] + nums2[(n+m+2)/2-1])*0.5
if not nums2:
return (nums1[(n+m+1)/2-1] + nums1[(n+m+2)/2-1])*0.5
def getKth(nums1, start1, end1, nums2, start2, end2, k):
"""
start1: 数组1第一个元素开始位置
end1: 数组1最后一个元素位置
start2: 数组2第一个元素开始位置
end2: 数组2最后一个元素位置
k:需要求的第K大元素
"""
if end1 < start1:
return nums2[start2+k-1]
if end2 < start2:
return nums1[start1+k-1]
if k == 1:
return min(nums1[start1], nums2[start2])
start1_temp = min(start1+k/2-1, len(nums1)-1)
start2_temp = min(start2+k/2-1, len(nums2)-1)
if nums2[start2_temp] > nums1[start1_temp]:
return getKth(nums1, start1_temp+1, end1, nums2, start2, end2, k-(start1_temp - start1+1))
else:
return getKth(nums1, start1, end1, nums2, start2_temp+1, end2, k-(start2_temp - start2+1))
return (getKth(nums1, 0, n-1, nums2, 0, m-1, (n+m+1)/2) + getKth(nums1, 0, n-1, nums2, 0, m-1, (n+m+2)/2))*0.5
总结
这题真是费了太大力气才调通了,主要原因是数组的下标太多了,一定要注意定义好每一个变量的循环不变量,另外这道题求中位数,更可怕的是有两种情况,这次直接通过(n+m+1)/2和(n+m+2)/2来完美融合两种情况。
- 合并奇偶问题: (n+m+1)/2和(n+m+2)/2
- 下标问题,明确循环不变量,end代表最后一个元素位置,start代表第一个元素位置,k代表第k大的数字位置。
- 越界问题:取数组下标和数组长度的最小值。
- 删除数组问题:修改start位置,同时更改k的大小。
- 空数组问题,递归函数要判断是否end>start。