Introduction

The original LeetCode problem can be found on this page.

Here in this page, we'll try to use binary search to solve it. Code sample is written in Go.

Merge Sort

At first glance, it's clear that we can use merge sort to merge the two arrays into one array, then finding the median would be easy. However, since we only need to find the median, there is no need to actually perform the merging, just finding the right indexes would suffice.

Here is the code where we try to find the indexes, with some explanation followed:

func findMedianSortedArrays(a []int, b []int) float64 {
	i, j := 0, 0
	for sizeOfFirstHalf := (len(a) + len(b) + 1) / 2; sizeOfFirstHalf > 0; sizeOfFirstHalf-- {
		if j == len(b) {
			i++
		} else if i == len(a) {
			j++
		} else if a[i] < b[j] {
			i++
		} else {
			j++
		}
	}

	var maxLeft int
	if j == 0 {
		maxLeft = a[i-1]
	} else if i == 0 {
		maxLeft = b[j-1]
	} else if a[i-1] > b[j-1] {
		maxLeft = a[i-1]
	} else {
		maxLeft = b[j-1]
	}
	if (len(a)+len(b))%2 == 1 {
		return float64(maxLeft)
	}

	var minRight int
	if j == len(b) {
		minRight = a[i]
	} else if i == len(a) {
		minRight = b[j]
	} else if a[i] < b[j] {
		minRight = a[i]
	} else {
		minRight = b[j]
	}
	return float64(maxLeft+minRight) / 2
}

• Imagine we do merge the two arrays into one sorted array, and we call that merged array.
• We divide the merged array into two equal sized halves, and put the median in the first half if the total number is odd.
• a[i:] and b[j:] belong to the second half of the merged array

Binary Search

Actually, there are some things about the indexes that we can utilize to perform a binary search:

• If we look at the for loop in previous code sample, we can see i + j == sizeOfFirstHalf after it, because we either advance i or j in each iteration. Thus, we can say j := sizeOfFirstHalf - i.
• a[i] >= b[j-1] because a[i] is in the second half while b[j-1] is in the first.
• The above is false for all a[:i], but is true for all a[i:], thus we can use binary search to find i.

Here is the code:

import "sort"

func findMedianSortedArrays(a []int, b []int) float64 {
	// make sure a is shorter so b[j-1] is always valid
	if len(a) > len(b) {
		return findMedianSortedArrays(b, a)
	}

	sizeOfFirstHalf := (len(a) + len(b) + 1) / 2
	i := sort.Search(len(a), func(i int) bool {
		j := sizeOfFirstHalf - i
		return a[i] >= b[j-1]
	})
	j := sizeOfFirstHalf - i

	var maxLeft int
	if j == 0 {
		maxLeft = a[i-1]
	} else if i == 0 {
		maxLeft = b[j-1]
	} else if a[i-1] > b[j-1] {
		maxLeft = a[i-1]
	} else {
		maxLeft = b[j-1]
	}
	if (len(a)+len(b))%2 == 1 {
		return float64(maxLeft)
	}

	var minRight int
	if j == len(b) {
		minRight = a[i]
	} else if i == len(a) {
		minRight = b[j]
	} else if a[i] < b[j] {
		minRight = a[i]
	} else {
		minRight = b[j]
	}
	return float64(maxLeft+minRight) / 2
}