↓Skip to main content
  1. Posts/

LeetCode-4 寻找两个正序数组的中位数

·1 min·
Table of Contents

LeetCode-4 寻找两个正序数组的中位数 #

Solution 1 #

分治. 目标是寻找有两个正序数组合并后的第 $k$ 个数 (这里 $k$ 从 1 开始记录). 考虑两个数组起始的坐标为 $idx_1$ 和 $idx_2$, 比较 $\text{nums1}[idx_1 + \lfloor \frac{k}{2}\rfloor - 1]$ 和 $\text{nums2}[idx_2 + \lfloor \frac{k}{2}\rfloor - 1]$ 的大小, 如果前者小于等于后者, 则说明 $\text{nums1}[idx_1 + \lfloor \frac{k}{2}\rfloor - 1]$ 及其之前的数都不可能是第 $k$ 个数, 可以排除掉. 否则, 排除 $\text{nums2}[idx_2 + \lfloor \frac{k}{2}\rfloor - 1]$ 及其之前的数. 更新 $idx$ 和 $k$ 的值, 递归调用直到 $k = 1$ 或者有一个数组已经遍历完.

代码如下:

class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        m, n = len(nums1), len(nums2)

        def f(i, j, k):
            if i == m:
                return nums2[j + k - 1]
            if j == n:
                return nums1[i + k - 1]
            if k == 1:
                return min(nums1[i], nums2[j])
            ni = min(m - 1, i + k // 2 - 1)
            nj = min(n - 1, j + k // 2 - 1)
            if nums1[ni] < nums2[nj]:
                return f(ni + 1, j, k - (ni - i + 1))
            else:
                return f(i, nj + 1, k - (nj - j + 1))

        if (m + n) & 1:
            return f(0, 0, (m + n + 1) // 2)
        return (f(0, 0, (m + n) // 2) + f(0, 0, (m + n) // 2 + 1)) / 2
↑