Find median of two sorted arrays

median.go

package main

func findMedianLog(a []int, b []int) float64 {
	n := len(a)
	m := len(b)

	meda, ia := median(a)
	medb, ib := median(b)

	if n == 0 {
		return medb
	}

	if m == 0 {
		return meda
	}

	if n == 1 && m == 1 {
		return (meda + medb) / 2
	}

	imin := min(ia, ib)
	if imin == 0 {
		imin = 1
	}

	if meda < medb {
		a = a[imin:n]
		b = b[0 : m-imin]
	} else {
		a = a[0 : n-imin]
		b = b[imin:m]
	}

	return findMedianLog(a, b)
}

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

	return b
}

func median(a []int) (float64, int) {
	n := len(a)
	if n == 0 {
		return 0, 0
	}

	if n%2 == 0 {
		return float64(a[n/2-1]+a[n/2]) / 2, n / 2
	}

	return float64(a[n/2]), n / 2
}

func findMedianLin(a []int, b []int) float64 {
	n := len(a)
	m := len(b)

	aIter := &iter{
		array: a,
	}

	bIter := &iter{
		array: b,
	}

	current := 0
	prev := 0

	for i := 0; i <= (n+m)/2; i++ {
		ai, aok := aIter.Current()
		bi, bok := bIter.Current()

		if aok && bok {
			if ai < bi {
				prev = current
				current = ai
				aIter.Next()
			} else {
				prev = current
				current = bi
				bIter.Next()
			}

			continue
		}

		if aok {
			prev = current
			current = ai
			aIter.Next()
			continue
		}

		if bok {
			prev = current
			current = bi
			bIter.Next()
			continue
		}
	}

	if (n+m)%2 == 0 {
		return float64(prev+current) / 2
	} else {
		return float64(current)
	}
}

type iter struct {
	array []int
	index int
}

func (i *iter) Next() bool {
	if i.index < len(i.array) {
		i.index++
		return true
	}

	return false
}

func (i *iter) Current() (int, bool) {
	if i.index < len(i.array) {
		return i.array[i.index], true
	}

	return 0, false
}

package main

import (
	"math"
	"testing"
)

func Test(t *testing.T) {
	tests := []struct {
		a []int
		b []int
		m float64
	}{
		{a: []int{}, b: []int{}, m: 0},
		{a: []int{}, b: []int{2}, m: 2},
		{a: []int{1}, b: []int{}, m: 1},
		{a: []int{1}, b: []int{2}, m: 1.5},
		{a: []int{1, 2}, b: []int{3, 4}, m: 2.5},
		{a: []int{1, 2, 3, 4, 5}, b: []int{1, 1, 2, 2, 3}, m: 2},
		{a: []int{1, 2, 3, 4}, b: []int{2, 3, 4, 5, 6, 7, 8}, m: 4},
	}

	var actual float64
	for _, test := range tests {
		actual = findMedianLin(test.a, test.b)
		if !almostEqual(actual, test.m) {
			t.Errorf("[linear] expected: %v, got: %v", test.m, actual)
		}

		actual = findMedianLog(test.a, test.b)
		if !almostEqual(actual, test.m) {
			t.Errorf("[log] expected: %v, got: %v", test.m, actual)
		}
	}
}

const float64EqualityThreshold = 1e-9

func almostEqual(a, b float64) bool {
	return math.Abs(a-b) <= float64EqualityThreshold
}