Go solution for https://leetcode.com/problems/median-of-two-sorted-arrays/

median_of_two_sorted_arrays.go

// https://leetcode.com/problems/median-of-two-sorted-arrays/
import (
    "fmt"
    "sort"
)

func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
    total := len(nums1) + len(nums2)
    mid1, mid2 := total / 2, total / 2
    if total % 2 == 0 {
        mid1--
    }
    // fmt.Printf("%d %d %d\n", mid1, mid2, total)
    if len(nums1) == 0 {
        return avg(nums2[mid1], nums2[mid2])
    } else if len(nums2) == 0 {
        return avg(nums1[mid1], nums1[mid2])
    }
    
    nums1, nums2 = order(nums1, nums2)
    if mid1 == mid2 {
        mid := findIdx(nums1, nums2, mid1)
        return avg(mid, mid)
    }
    return avg(findIdx(nums1, nums2, mid1), findIdx(nums1, nums2, mid2))   
}

func order(x, y []int) ([]int, []int) {
    if x[0] < y[0] {
        return x, y
    } else {
        return y, x
    }
}

func avg(a, b int) float64 {
    return float64(a + b) / 2.0
}

// x[0] is guaranteed <= y[0]
func findIdx(x, y []int, idx int) int {
    _, b, c, _ := x[0], x[len(x)-1], y[0], y[len(y)-1]
    overlap_x, overlap_y := sort.SearchInts(x, c), sort.SearchInts(y, b)
    // fmt.Printf("%d %d %d %d %d\n", b, c, overlap_x, overlap_y, idx)
    if idx <= overlap_x - 1 {
      return x[idx]
    } else if idx >= len(x) + overlap_y {
      return y[idx - len(x)]
    } else {
      return linearSearch(x[overlap_x:], y[:overlap_y], idx - overlap_x)
    }
}

func linearSearch(x, y []int, idx int) int {
    // fmt.Printf("%+v %+v %d\n", x, y, idx)
    xPos, yPos := 0, 0
    var val int
    if idx > len(x) + len(y) {
        panic("impossible!")
    }
    for idx >= 0 {
        if yPos >= len(y) || (xPos < len(x) && x[xPos] < y[yPos]) {
            // fmt.Printf("%d %d\n", xPos, len(x))
            val = x[xPos]
            xPos++          
        } else {
            val = y[yPos]
            yPos++
        }
        idx--
    }
    return val
}