leetcode-go/hard/4.go

package main

import "fmt"

// len(nums1) + len(nums2) != 0
func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
	// Use i, j to divide A, B into 2 parts, like
	// [A(0), ... , A(i-1), | A(i), ... , A(m)] and
	// [B(0), ... , B(j-1), | B(j), ... , B(n)],
	// then the left part of A and B is left, the right
	// part is right, when max(left) <= min(right) &&
	// len(left) == len(right), we can find the median by
	// median = [max(left) + min(right)] / 2
	m, n := len(nums1), len(nums2)
	// if m < n, i ~ [0, m], j = (m+n+1)/2 - i
	if m > n { // if m > n, swap nums1, nums2
		nums1, nums2, m, n = nums2, nums1, n, m
	}
	ibeg, iend, halfLen := 0, m, (m+n+1)/2
	var i, j int
	for ibeg <= iend {
		i = (ibeg + iend) / 2
		j = halfLen - i
		if i > 0 && nums1[i-1] > nums2[j] { // A[i-1]'s too big, ie. i's too big
			iend = i - 1
		} else if i < m && nums2[j-1] > nums1[i] { // i's too small
			ibeg = i + 1
		} else {
			// i, j is found (A[i-1] <= B[j] && B[j-1] <= A[i])
			var maxOfLeft, minOfRight int
			if i == 0 { // find max of left
				maxOfLeft = nums2[j-1]
			} else if j == 0 {
				maxOfLeft = nums1[i-1]
			} else {
				maxOfLeft = nums1[i-1]
				if nums2[j-1] > maxOfLeft {
					maxOfLeft = nums2[j-1]
				}
			}
			if (m+n)%2 != 0 { // if odd, return max of left as median (!!!)
				return float64(maxOfLeft)
			}
			if i == m { // find min of right
				minOfRight = nums2[j]
			} else if j == n {
				minOfRight = nums1[i]
			} else {
				minOfRight = nums1[i]
				if nums2[j] < nums1[i] {
					minOfRight = nums2[j]
				}
			}
			return float64(maxOfLeft+minOfRight) / 2
		}
	}
	return -1 // not found, useless
}

func main() {
	nums1 := []int{1, 3}
	nums2 := []int{2}
	fmt.Println(findMedianSortedArrays(nums1, nums2))
	nums3 := []int{}
	nums4 := []int{1}
	fmt.Println(findMedianSortedArrays(nums3, nums4))
}