题目描述

给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的 中位数 。
算法的时间复杂度应该为 O(log (m+n)) 。

题目分析

设总数量为 sum，问题可以转化为
– sum 是奇数，寻找第 sum / 2 + 1 个数字
– sum 是偶数，寻找第 sum / 2 和 sum / 2 + 1 个数字，再取平均值

因此问题转化为在 2 个有序数组中寻找第 k 小的数字。
分别取 2 个数组中第 k / 2 个数字 x1 和 x2：
– x1 < x2： x1 至多是第 k – 1 个数字，x1 及左侧都可以丢弃掉，递归处理 x1 后面的数字和 nums2
– x1 >= x2： x2 至多是第 k – 1 个数字，x2 及左侧都可以丢弃掉，递归处理 x2 后面的数字和 nums1

这样每次可以丢弃 k / 2 个数字，递归出口是 num1 和 num2 中有一个为空，或 k = 1（两数组的第一个数字取最小值）。
k 和 sum 成正比，因此时间复杂度 O(log(m + n)).

Java

public double findMedianSortedArrays(int[] nums1, int[] nums2) {
    int sum = nums1.length + nums2.length;
    if (sum % 2 == 0) {
        int ans = findK(nums1, 0, nums2, 0, sum / 2) + findK(nums1, 0, nums2, 0, sum / 2 + 1);
        return ans / 2.0;
    } else {
        return findK(nums1, 0, nums2, 0, sum / 2 + 1);
    }
}

private int findK(int[] nums1, int start1, int[] nums2, int start2, int k) {
    if (start1 == nums1.length) {
        return nums2[start2 + k - 1];
    }
    if (start2 == nums2.length) {
        return nums1[start1 + k - 1];
    }
    if (k == 1) {
        return Math.min(nums1[start1], nums2[start2]);
    }

    int mid = k / 2 - 1;
    int index1 = nums1.length > start1 + mid ? start1 + mid : nums1.length - 1;
    int index2 = nums2.length > start2 + mid ? start2 + mid : nums2.length - 1;
    if (nums1[index1] < nums2[index2]) {
        return findK(nums1, index1 + 1, nums2, start2, k - (index1 - start1 + 1));
    } else {
        return findK(nums1, start1, nums2, index2 + 1, k - (index2 - start2 + 1));
    }
}

Kotlin

fun findMedianSortedArrays(nums1: IntArray, nums2: IntArray): Double {
    val sum = nums1.size + nums2.size
    if (sum % 2 == 0) {
        val ans = findK(nums1, 0, nums2, 0, sum / 2) + findK(nums1, 0, nums2, 0, sum / 2 + 1)
        return ans / 2.0
    } else {
        return findK(nums1, 0, nums2, 0, sum / 2 + 1).toDouble()
    }
}

private fun findK(nums1: IntArray, start1: Int, nums2: IntArray, start2: Int, k: Int): Int {
    if (start1 == nums1.size) {
        return nums2[start2 + k - 1]
    }
    if (start2 == nums2.size) {
        return nums1[start1 + k - 1]
    }
    if (k == 1) {
        return min(nums1[start1], nums2[start2])
    }
    val mid = k / 2 - 1
    val index1 = if (nums1.size > mid + start1) mid + start1 else nums1.size - 1
    val index2 = if (nums2.size > mid + start2) mid + start2 else nums2.size - 1
    return if (nums1[index1] < nums2[index2]) {
        findK(nums1, index1 + 1, nums2, start2, k - (index1 - start1 + 1))
    } else {
        findK(nums1, start1, nums2, index2 + 1, k - (index2 - start2 + 1))
    }
}