leetcode-go/leetcode/medium/658.find-k-closest-elements.go

func findClosestElements(arr []int, k int, x int) (res []int) {
	n := len(arr)
	beg, end := 0, n-1
	for beg <= end {
		mid := beg + (end-beg)/2
		if arr[mid] < x {
			beg = mid + 1
		} else {
			end = mid - 1
		}
	}
	if beg == n {
		return arr[n-k:]
	} else if arr[beg] == x {
		res = append(res, x)
		beg, end = beg-1, beg+1
		k--
	} else {
		beg, end = beg-1, beg
	}
	for ; 0 < k; k-- {
		if beg < 0 {
			res = append(res, arr[end])
			end++
		} else if n <= end {
			res = append(res, arr[beg])
			beg--
		} else {
			if x-arr[beg] <= arr[end]-x {
				res = append(res, arr[beg])
				beg--
			} else {
				res = append(res, arr[end])
				end++
			}
		}
	}
	sort.Ints(res)
	return res
}