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

题目链接

还有两道相关题。下面代码中的一部分可以拿来解这两道题。

1. NC36 在两个长度相等的排序数组中找到上中位数
2. CD82 在两个排序数组中找到第k小的数

核心思路：根据有序数组的性质，不断抛弃不存在所需数据的区间。
这确实是一道难题，需要细致分析多种情况。最好的办法就是在纸上分情况推演。

// 入口函数
func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
    l1 := len(nums1)
    l2 := len(nums2)
    if l1 == 0 && l2 == 0 {
        panic("nums1 and nums2 are all empty")
    }
    if l1 == 0 {
        return getMedian(nums2)
    }
    if l2 == 0 {
        return getMedian(nums1)
    }
    // l1 > 0 && l2 > 0
    l := l1+l2
    if l&1 != 0 {
        return float64(getKthSmallest(nums1, nums2, l>>1+1))
    }
    return (float64(getKthSmallest(nums1, nums2, l>>1))+float64(getKthSmallest(nums1, nums2, l>>1+1)))/2
}

// 找到两个有序数组的第k小的数
// k >= 1 && k <= len(a1)+len(a2)
func getKthSmallest(a1, a2 []int, k int) int {
    var (
        l1 = len(a1)
        l2 = len(a2)
    )
    if k < 1 || k > l1+l2 {
        panic("k is invalid")
    }
    var (
        sl = l1 // 短数组的长度
        sa = a1 // 短数组
        ll = l2 // 长数组的长度
        la = a2 // 长数组
    )
    if l1 > l2 {
        sl = l2
        sa = a2
        ll = l1
        la = a1
    }
    if k <= sl {
        return getUpMedian(a1[:k], a2[:k])
    }
    if k > ll {
        li := k-sl-1 // 长数组的索引
        if la[li] >= sa[sl-1] {
            return la[li]
        }
        si := k-ll-1
        if sa[si] >= la[ll-1] {
            return sa[si]
        }
        return getUpMedian(sa[si+1:], la[li+1:])
    }
    // k > sl && k <= ll
    li := k-sl-1 // 长数组索引
    if la[li] >= sa[sl-1] {
        return la[li]
    }
    return getUpMedian(sa, la[li+1:k])
}

// 找到两个长度相等的有序数组的上中位数
func getUpMedian(a1, a2 []int) int {
    var (
        l1 = len(a1)
        l2 = len(a2)
    )
    if l1 != l2 {
        panic("len(a1) != len(a2)")
    }
    var (
        b1 int
        e1 = l1-1
        b2 int
        e2 = l2-1
    )
    for b1 < e1 && b2 < e2 {
        m1 := b1+(e1-b1)>>1
        m2 := b2+(e2-b2)>>1
        l := e1-b1+1
        if a1[m1] == a2[m2] {
            return a1[m1]
        }
        if a1[m1] > a2[m2] {
            if l&1 != 0 {
                e1 = m1
                b2 = m2
            } else {
                e1 = m1
                b2 = m2+1
            }
        } else {
            if l&1 != 0 {
                b1 = m1
                e2 = m2
            } else {
                b1 = m1+1
                e2 = m2
            }
        }
    }
    return min(a1[b1], a2[b2])
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

// 计算一个有序数组的中位数
func getMedian(nums []int) float64  {
    l := len(nums)
    if l == 0 {
        panic("nums is empty")
    }
    if l&1 != 0 {
        return float64(nums[l>>1])
    }
    return (float64(nums[l>>1-1])+float64(nums[l>>1]))/2
}
复制代码

时间复杂度$O(log(min(m,n)))$，空间复杂度$O(1)$。