Top-K Elements

A heap (priority queue) efficiently maintains the min or max of a dynamic collection. Go’s container/heap provides the interface.

Use when: you need the k largest/smallest elements, merge k sorted lists, or schedule by priority.

To use Go’s heap, implement the heap.Interface on a slice type. This is a min-heap of ints.

Find the k largest elements using a min-heap of size k. Counterintuitive: we use a MIN-heap because we want to efficiently drop the smallest of our k candidates.

If heap exceeds size k, remove the smallest. What remains are the k largest seen so far.

Last Stone Weight: repeatedly smash the two heaviest stones. This is a natural max-heap problem. We negate values to simulate a max-heap with Go’s min-heap.

Trace: topK([3, 1, 5, 12, 2, 11], 3)

Push 3: heap [3] Push 1: heap [1, 3] Push 5: heap [1, 3, 5] Push 12: heap [1, 3, 5, 12], pop 1 -> [3, 5, 12] Push 2: heap [2, 3, 5, 12], pop 2 -> [3, 5, 12] Push 11: heap [3, 5, 11, 12], pop 3 -> [5, 11, 12] Result: [5, 11, 12]

type MinHeap []int

func (h MinHeap) Len() int            { return len(h) }
func (h MinHeap) Less(i, j int) bool  { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}
func topK(nums []int, k int) []int {
	h := &MinHeap{}
	heap.Init(h)

	for _, n := range nums {
		heap.Push(h, n)
		if h.Len() > k {
			heap.Pop(h)
		}
	}

	return []int(*h)
}
func lastStoneWeight(stones []int) int {
	h := &MinHeap{}
	for _, s := range stones {
		heap.Push(h, -s)
	}

	for h.Len() > 1 {
		first := -heap.Pop(h).(int)
		second := -heap.Pop(h).(int)

		if first != second {
			heap.Push(h, -(first - second))
		}
	}

	if h.Len() == 0 {
		return 0
	}
	return -heap.Pop(h).(int)
}
func main() {
	fmt.Println("Top 3:", topK([]int{3, 1, 5, 12, 2, 11}, 3))
	fmt.Println("Last stone:", lastStoneWeight([]int{2, 7, 4, 1, 8, 1}))
}