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

题目

给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的 中位数 。

算法的时间复杂度应该为 O(log (m+n)) 。

难度：🌟🌟🌟

示例 1：

输入：nums1 = [1,3], nums2 = [2]
输出：2.00000
解释：合并数组 = [1,2,3] ，中位数 2

示例 2：

输入：nums1 = [1,2], nums2 = [3,4]
输出：2.50000
解释：合并数组 = [1,2,3,4] ，中位数 (2 + 3) / 2 = 2.5

提示：

• nums1.length == m
• nums2.length == n
• 0 <= m <= 1000
• 0 <= n <= 1000
• 1 <= m + n <= 2000
• -10⁶ <= nums1[i], nums2[i] <= 10⁶

题解

借鉴归并排序中合并两个有序数组的思路，在合并时进行计数，当下标等于中位数下标时，记录此数，最后计算中位数。

代码实现

func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
	m, n := len(nums1), len(nums2)
	var isOdd bool // 是否是偶数
	medianNumIdx := 0
	if (m+n)%2 == 0 { // 偶数
		isOdd = true
		medianNumIdx = (m+n)/2 - 1
	} else { // 奇数
		isOdd = false
		medianNumIdx = (m + n) / 2
	}
	count := 0
	i, j := 0, 0
	var medianNum []int
	for i < m && j < n {
		var median int
		if nums1[i] < nums2[j] {
			median = nums1[i]
			i++
		} else {
			median = nums2[j]
			j++
		}
		if count == medianNumIdx {
			medianNum = append(medianNum, median)
		}
		// 偶数需要额外处理
		if isOdd && count == medianNumIdx+1 {
			medianNum = append(medianNum, median)
		}
		count++
	}

	for i < m {
		median := nums1[i]
		if count == medianNumIdx {
			medianNum = append(medianNum, median)
		}
		// 偶数需要额外处理
		if isOdd && count == medianNumIdx+1 {
			medianNum = append(medianNum, median)
		}
		count++
		i++
	}

	for j < n {
		median := nums2[j]
		if count == medianNumIdx {
			medianNum = append(medianNum, median)
		}
		// 偶数需要额外处理
		if isOdd && count == medianNumIdx+1 {
			medianNum = append(medianNum, median)
		}
		count++
		j++
	}

	if isOdd {
		return float64(medianNum[0]+medianNum[1]) / 2
	}
	return float64(medianNum[0])
}