leetcode(4) Median of Two Sorted Arrays

problem

There are two sorted arrays nums1 and nums2 of size m and n
respectively.

Find the median of the two sorted arrays. The overall run time
complexity should be O(log (m+n)).

Example 1: nums1 = [1, 3] nums2 = [2]

The median is 2.0
Example 2: nums1 = [1, 2] nums2 = [3, 4]

solution

这个问题最初的一个想法就是设两个指针，从小到大分别遍历两个数组，直到遍历$(m+n)/2$个元素，但这种算法的时间复杂度为$O(m+n)$，在这篇博客中看到了一个不错的方法（感觉优秀的算法都是一个思路的转变），例如在two sum中就是把判断转换为查找问题，而在这个问题中则是把找中位数转换为寻找第k小元素的问题，这两种思路区别主要在于中位数要考虑两侧元素的比较，而寻找第k小元素只需要从一端开始进行比较，而又因为原始序列有序因此可以进行二分优化。具体解答如下：
寻找两个升序排列数组A，B中的第k小的元素：
首先假设两个数组的长度都大于k/2，我们首先比较A[k/2-1]和B[k/2-1]两个元素，这两个元素分别代表两个数组中第k/2小的元素，这两个元素比较有三种情况：

1. A[k/2-1]< B[k/2-1]，这时可以得出结论A[k/2-1]及之前的元素一定小于第k小元素
2. A[k/2-1]>B[k/2-1]，同1类似，B[k/2-1]及之前的元素一定小于第k小元素
3. A[k/2-1]=B[k/2-1]，A[k/2-1]就是第k小的元素

这样我们只需递归查找，此外还要考虑几个边界条件：

1. 如果A或者B为空，则直接返回B[k-1]或者A[k-1]；
2. 如果k为1，我们只需要返回A[0]和B[0]中的较小值；
3. 如果A[k/2-1]=B[k/2-1]，返回其中一个；

寻找第k小元素具体代码如下：

def findKth(a, b, m, n, k):
    if m > n:
        return findKth(b, a, n, m, k)
    if m==0:
        return b[k-1]
    if k==1:
        return min(a[0], b[0])

    pa = min(k//2, m)
    pb = k-pa
    if a[pa-1]<b[pb-1]:
        return findKth(a[pa:], b, m-pa, n, k-pa)
    elif a[pa-1]>b[pb-1]:
        return findKth(a, b[pb:], m, n-pb, k-pb)
    else:
        return a[pa-1]

想要找中位数只要按奇偶数分别找对应的元素即可，时间复杂度为$O(log(m+n))$，空间复杂度为$O(m+n)$

总结

一个朴素的想法就是对两个数组从头类似于MergeSort一样进行排序，找到中位数的位置，但是这样的算法是一个一个的排除不符合条件的元素的，而原数组是有序的，我们可以把二分查找的思想引入进来，即一次删除一半不符合条件的元素，最终算法的时间复杂度为$O(log(m+n))$

其实这个想法一开始我也想到了，只不过不知道如何写成代码，没有想到用递归的方法实现，还要多加练习。

扩展

如果传入的有序数组是多个（如3个）的话算法应该怎样调整？

答：如给出三个有序的数组A,B,C，求整个元素的中位数的话，假设三个数组的长度都大于k/3，那么比较A[k/3-1], B[k/3-1], C[k/3-1]，有三种情况：

1. 这三个数都不相同，设最大的是C[k/3-1]，那么A[k/3-1], B[k/3-1]一定小于第k个数，所以删除这两个序列前k个元素。
2. A[k/3-1]< B[k/3-1]= C[k/3-1]，那么删除A中前k个元素，递归
3. 这三个数相同，那么第k个数的值就是A[k/3-1]

思想和上面算法类似，就是如何正确删除最多的元素，
时间复杂度分析：
时间复杂度很简单仍然为$O(\log(m+n+p))$，下面定量比较一下两个数组和三个数组的时间复杂度。
递归深度最坏为$\log_{1.5}(m+n+p)$、最好为$\log_3(m+n+p)$
每次比较3次，总时间$\approx$3倍递归深度
而两个数组的话因为每次只需比较一次，因此总时间$\approx$递归深度