kClosest.go

package heap

func KClosest(points [][]int, k int) [][]int {
	if len(points) == 0 {
		return [][]int{}
	}
	var heap [][]int

	for _, pair := range points {
		x, y := pair[0], pair[1]
		dist := x*x + y*y
		heap = append(heap, []int{dist, x, y})
	}

	buildHeap(heap)

	var results [][]int
	for k > 0 {
		root := pop(&heap)
		x, y := root[1], root[2]
		results = append(results, []int{x, y})
		k--
	}

	return results
}

func buildHeap(heap [][]int) {
	n := len(heap)
	startIdx := n / 2 - 1

	for i := startIdx; i >= 0; i-- {
		heapify(&heap, i)
	}
}

func heapify(heap *[][]int, i int) {
	smallest := i
	lChild := 2*i + 1
	rChild := 2*i + 2

	if lChild < len(*heap) && (*heap)[lChild][0] < (*heap)[smallest][0] {
		smallest = lChild
	}
	if rChild < len(*heap) && (*heap)[rChild][0] < (*heap)[smallest][0] {
		smallest = rChild
	}

	if smallest != i {
		(*heap)[i], (*heap)[smallest] = (*heap)[smallest], (*heap)[i]
		heapify(heap, smallest)
	}
}

func pop(heap *[][]int) []int {
	root := (*heap)[0]

	n := len(*heap)
	lastElement := (*heap)[n - 1]
	(*heap)[0] = lastElement
	*heap = (*heap)[:n-1]
	heapify(heap, 0)

	return root
}