K Closest Points to Origin

heap.go

package main

import (
	"fmt"
	"strings"
)

func main() {
	p := [][]int{
		[]int{17, -45},
		[]int{-72, -76},
		[]int{21, 99},
		[]int{-25, -19},
		[]int{-48, 86},
		[]int{86, -47},
		[]int{59, -66},
		[]int{-38, -16},
		[]int{47, -44},
		[]int{28, 96},
		[]int{92, -64},
		[]int{-62, -35},
		[]int{-63, -87},
		[]int{-83, 95},
		[]int{25, -89},
		[]int{30, -40},
		[]int{-75, 93},
		[]int{47, -21},
	}
	fmt.Println(kClosest(p, 9))
}

func kClosest(points [][]int, K int) [][]int {
	h := NewHeap(K)
	for _, c := range points {
		p := NewPoint(c)
		h.Add(p)
	}
	r := make([][]int, K)
	for i, p := range h.points {
		r[i] = p.coords
	}
	return r
}

type Heap struct {
	capacity int
	filled   int
	points   []*Point
}

func (h *Heap) String() string {
	return "\n" + h.StringTree(0, 0)
}
func (h *Heap) StringTree(i, d int) string {
	r := fmt.Sprintf("|-%v %v\n", strings.Repeat("-", 4*d), h.points[i])
	if 2*i+1 < h.capacity {
		r += h.StringTree(2*i+1, d+1)
	}
	if 2*i+2 < h.capacity {
		r += h.StringTree(2*i+2, d+1)
	}
	return r
}

func NewHeap(K int) *Heap {
	return &Heap{K, 0, make([]*Point, K)}
}

func (h *Heap) Less(a, b int) bool {
	return h.points[a].Less(h.points[b])
}

func (h *Heap) Swap(a, b int) {
	h.points[a], h.points[b] = h.points[b], h.points[a]
}

/* This is a custom method to keep the heap length = K
* If there's space it does a simple heap.Push
* Otherwise if the new point is closer than the furthest (heap root)
* It does a heap.Pop and then a heap.Push with the new point.
 */
func (h *Heap) Add(p *Point) {
	fmt.Print("New point: ", p)
	if h.filled < h.capacity {
		h.Push(p)
		fmt.Println(" -> Pushed, heap:", h.points, h)
	} else if p.Less(h.points[0]) {
		h.Pop()
		h.Push(p)
		fmt.Println(" -> Popped, heap:", h.points, h)
	} else {
		fmt.Println(" -> Ignored\n")
	}
}

func (h *Heap) Push(p *Point) {
	// Assume h.filled < h.capacity
	i := h.filled
	h.points[i] = p
	h.filled++
	// Bubble up
	for i > 0 && h.Less((i-1)/2, i) {
		h.Swap(i, (i-1)/2)
		i = (i - 1) / 2
	}
}

func (h *Heap) Pop() {
	h.filled--
	h.points[0] = h.points[h.filled]
	i := 0
	// Bubble down
	for 2*i+1 < h.capacity {
		l, r := 2*i+1, 2*i+2
		// There may not be a right child
		if r >= h.capacity {
			if h.Less(l, i) {
				// Everything is in the right place stop
				return
			}
			h.Swap(i, l)
			// We've moved to the end of the heap so we can stop
			return
		}
		if h.Less(l, i) && h.Less(r, i) {
			// The current node is greater than both children so we can stop
			return
		}
		// Otherwise swap with larger child and continue
		if h.Less(l, r) {
			h.Swap(i, r)
			i = r
		} else {
			h.Swap(i, l)
			i = l
		}
	}
}

type Point struct {
	coords []int
	sqdist int
}

func (p *Point) Less(q *Point) bool {
	return p.sqdist < q.sqdist
}

func NewPoint(p []int) *Point {
	s := p[0]*p[0] + p[1]*p[1]
	return &Point{p, s}
}

func (p *Point) String() string {
	return fmt.Sprintf("%d(%d,%d)", p.sqdist, p.coords[0], p.coords[1])
}