[LeetCode] Median of Two Sorted Arrays

题意

给你俩排序过的数组 A 和 B，现在把俩数组混合在一起，求混合后数组的中位数。

要求：时间复杂度是 $O(log (m+n))$ 。

题解

大力出奇迹

反正 LeetCode 不怎么卡时间……

那就当成模拟题来做。

代码

class Solution(object):
    def findMedianSortedArrays(self, nums1, nums2):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float
        """
        t = nums1
        t.extend(nums2)
        t.sort()
        length = len(t)
        t_m = length >> 1
        if length & 1:
            return t[t_m]
        else:
            return (t[t_m - 1] + t[t_m]) / 2.0

结果

2080 / 2080 test cases passed.
Status: Accepted
Runtime: 124 ms

2080 个测试点…… 我的小心肝……

递归版

出题人肯定是要让我们用二分的。都告诉你时间复杂度要求了啊！

想法是好的：如果能求俩数组中的任意第 k 小数，不就完事儿了么？

首先，在俩数组中分别取中位数，分别是第 i 个和第 j 个。i 和 j 将 A 和 B 数组分成了四个部分：A 数组中比 A[i] 小的、大的；B 数组中比 B[j] 小的、大的：

aaaaaaa i AAAAAAA
bbbbbbb j BBBBBBB

然后，看 i+j 和 k 的大小。如果 i+j<k，那么说明，i 和 j 整体取小了。这时需要舍去 a 和 b 中较小的一段。“较小的” 怎么定义呢？看 A[i] 和 B[j] 谁更小。反之，说明 i 和 j 整体取大了，舍去 A 和 B 中较大的一段。

思路很好啊，我没写出来…… 囧…… 原因是想写一个非递归的版本，然后一直在边界条件上绕不过来弯。

几经波折，最后写出了一个这样的版本。

代码

class Solution(object):
    def findKthBig(self, a, b, k):
        a_len = len(a)
        b_len = len(b)
        if a_len == 0:
            return b[k - 1]
        elif b_len == 0:
            return a[k - 1]
        elif k == 1:
            return min(a[0], b[0])

        k_m = k >> 1
        a_m = min(a_len, k_m)
        b_m = min(b_len, k_m)
        if a[a_m - 1] < b[b_m - 1]:
            return self.findKthBig(a[a_m:], b, k - a_m)
        else:
            return self.findKthBig(a, b[b_m:], k - b_m)

    def findMedianSortedArrays(self, nums1, nums2):
        len1 = len(nums1)
        len2 = len(nums2)
        mid = (len1 + len2) >> 1
        ans = 0

        if (len1 + len2) & 1:
            ans = self.findKthBig(nums1, nums2, mid + 1)
        else:
            ans = (self.findKthBig(nums1, nums2, mid) +
                   self.findKthBig(nums1, nums2, mid + 1)
                   ) / 2.0

        return ans

结果

2080 / 2080 test cases passed.
Status: Accepted
Runtime: 138 ms

非递归版

由上面的递归版本改写成了下面的非递归的版本：

代码

class Solution(object):
    def findKthBig(self, a, b, k):
        a_len = len(a)
        b_len = len(b)
        a_l = 0
        b_l = 0
        while True:
            if a_len == 0:
                return b[b_l + k - 1]
            elif b_len == 0:
                return a[a_l + k - 1]
            elif k == 1:
                return min(a[a_l], b[b_l])

            k_m = k >> 1

            a_len_mid = min(a_len, k_m)
            a_m = a_l + a_len_mid
            a_k = k - a_len_mid

            b_len_mid = min(b_len, k_m)
            b_m = b_l + b_len_mid
            b_k = k - b_len_mid

            if a[a_m - 1] < b[b_m - 1]:
                a_l = a_m
                a_len -= a_len_mid
                k = a_k
            else:
                b_l = b_m
                b_len -= b_len_mid
                k = b_k

    def findMedianSortedArrays(self, nums1, nums2):
        len1 = len(nums1)
        len2 = len(nums2)
        mid = (len1 + len2) >> 1
        ans = 0

        if (len1 + len2) & 1:
            ans = self.findKthBig(nums1, nums2, mid + 1)
        else:
            ans = (self.findKthBig(nums1, nums2, mid) +
                   self.findKthBig(nums1, nums2, mid + 1)
                   ) / 2.0

        return ans

结果

2080 / 2080 test cases passed.
Status: Accepted
Runtime: 179 ms

为什么时间反而长了？

桩程序

if __name__ == '__main__':
    cases = [
        {'nums1': [1, 2, 3, 4, 5, 6, 7, 8, 9], 'nums2': [4, 6, 8, 10, 12], 'ans': 6},
        {'nums1': [1, 2], 'nums2': [3, 4], 'ans': 2.5},
        {'nums1': [1, 3], 'nums2': [2], 'ans': 2.0},
        {'nums1': [1, 2, 3, 4, 5], 'nums2': [5, 6, 8, 10, 12], 'ans': 5},
        {'nums1': [], 'nums2': [3, 4], 'ans': 3.5},
        {'nums1': [1, 2], 'nums2': [], 'ans': 1.5},
    ]

    for item in cases:
        ans = (Solution()).findMedianSortedArrays(
            nums1=item['nums1'],
            nums2=item['nums2']
        )

        print(str(ans == item['ans']) + "\t=>\t" + str(ans) + "\t:\t" + str(item['ans']))