寻找两个有序数组的中位数

发布时间 2023-04-07 00:22:04作者: 叶天云

题目链接

解题思路

由于要求时间复杂度O(log(m+n)),所以使用二分法
寻找两个数组合并后的中位数,其实就是找第(m+n)/2或者(m+n)/2+1小的数,
假设这个数是k(k是整数, k/2是整除),比较nums1[k/2 - 1]和nums2[k/2 - 1]的值
图片1
1.如果前者小于后者, 比nums1[k/2 - 1]小的数只有nums1[0]到nums1[k/2 - 1]和nums2[0]到nums2[k/2 - 1]个,加起来就是k - 2个, 所以nums1[0]到nums1[k/2 - 1]绝对不是第k个值,直接排除掉这些数,更新k = k - k/2继续下个循环
图片3
2.如果前者大于后者,那就排除nums2[0]到nums2[k/2 -1],更新k = k - k/2,继续下个循环
图片2
3.如果两者相等,那就归入第一种情况
特殊结果
1.如果有一个数组为空,那就返回另一个数组的第k个元素
图片5
2.如果k=1,那就返回两个数组的首值最小数
3.如果越界,那就取数组里的最后一个
图片4

代码

func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
 totalLength := len(nums1) + len(nums2)
 // 中位数的奇数和偶数计算公式
 if totalLength%2 == 1 {
  midIndex := totalLength / 2
  return float64(getKthElement(nums1, nums2, midIndex+1))
 } else {
  midIndex1, midIndex2 := totalLength/2-1, totalLength/2
  return float64(getKthElement(nums1, nums2, midIndex1+1)+getKthElement(nums1, nums2, midIndex2+1)) / 2.0
 }
}

func getKthElement(nums1, nums2 []int, k int) int {
 index1, index2 := 0, 0
 for {
  // 数组1是空的
  if index1 == len(nums1) {
   // 直接取数组2的值
   return nums2[index2+k-1]
  }
  // 数组2是空的
  if index2 == len(nums2) {
   //直接取数组1的值
   return nums1[index1+k-1]
  }
  // k = 1的情况
  if k == 1 {
   return min(nums1[index1], nums2[index2])
  }
  half := k / 2
  // 取nums1[k/2 -1] nums2[k/2 -1]
  // 并且判断是否越界,如果越界取最后一个元素
  newIndex1 := min(index1+half, len(nums1)) - 1
  newIndex2 := min(index2+half, len(nums2)) - 1
  pivot1, pivot2 := nums1[newIndex1], nums2[newIndex2]
  // 根据结果更新k的值和下次循环的指针下标
  if pivot1 <= pivot2 {
   k -= (newIndex1 - index1 + 1)
   index1 = newIndex1 + 1
  } else {
   k -= (newIndex2 - index2 + 1)
   index2 = newIndex2 + 1
  }
 }
}

func min(x, y int) int {
 if x < y {
  return x
 }
 return y
}

测试结果

func main() {
 nums1 := []int{1, 2, 3}
 nums2 := []int{3, 3, 3}

 res1 := findMedianSortedArrays(nums1, nums2)
 fmt.Println(res1)

 nums3 := []int{1, 2}
 nums4 := []int{3, 4}
 res2 := findMedianSortedArrays(nums3, nums4)
 fmt.Println(res2)
}