Heap / Top-K

// IntHeap is a MIN-heap of ints. container/heap needs five methods:
// the three from sort.Interface plus Push and Pop (which only append / trim the slice).
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] } // flip to > for a max-heap
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *IntHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *IntHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

// Usage: h := &IntHeap{}; heap.Push(h, 5); smallest := (*h)[0]; heap.Pop(h).(int)

// FindKthLargest: keep a min-heap of the k largest values seen so far.
// The root is the smallest of those k, which is exactly the kth largest overall.
func FindKthLargest(nums []int, k int) int {
	h := &IntHeap{}
	for _, v := range nums {
		heap.Push(h, v)
		if h.Len() > k {
			heap.Pop(h) // evict the smallest: it can no longer be in the top k
		}
	}
	return (*h)[0]
}

type freqItem struct{ val, count int }
type freqHeap []freqItem

func (h freqHeap) Len() int { return len(h) }
func (h freqHeap) Less(i, j int) bool {
	if h[i].count != h[j].count {
		return h[i].count < h[j].count
	}
	return h[i].val > h[j].val // deterministic tie-break
}
func (h freqHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *freqHeap) Push(x any)   { *h = append(*h, x.(freqItem)) }
func (h *freqHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

// TopKFrequent: count with a map, then keep a size-k min-heap by count.
// Result is ordered from most to least frequent.
func TopKFrequent(nums []int, k int) []int {
	count := map[int]int{}
	for _, v := range nums {
		count[v]++
	}
	h := &freqHeap{}
	for val, c := range count {
		heap.Push(h, freqItem{val, c})
		if h.Len() > k {
			heap.Pop(h)
		}
	}
	res := make([]int, h.Len())
	for i := len(res) - 1; i >= 0; i-- {
		res[i] = heap.Pop(h).(freqItem).val
	}
	return res
}

// ListNode is the singly linked list node.
type ListNode struct {
	Val  int
	Next *ListNode
}

type nodeHeap []*ListNode

func (h nodeHeap) Len() int           { return len(h) }
func (h nodeHeap) Less(i, j int) bool { return h[i].Val < h[j].Val }
func (h nodeHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *nodeHeap) Push(x any)        { *h = append(*h, x.(*ListNode)) }
func (h *nodeHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

// MergeKLists: the heap always holds the current head of every list;
// pop the smallest, append it, and push that node's successor.
func MergeKLists(lists []*ListNode) *ListNode {
	h := &nodeHeap{}
	for _, l := range lists {
		if l != nil {
			heap.Push(h, l)
		}
	}
	dummy := &ListNode{}
	tail := dummy
	for h.Len() > 0 {
		n := heap.Pop(h).(*ListNode)
		tail.Next = n
		tail = n
		if n.Next != nil {
			heap.Push(h, n.Next)
		}
	}
	return dummy.Next
}

type maxHeap []int

func (h maxHeap) Len() int           { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h maxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *maxHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

// MedianFinder: a max-heap holds the smaller half, a min-heap the larger half.
// Invariant: len(lo) == len(hi) or len(lo) == len(hi)+1, and every lo value <= every hi value.
type MedianFinder struct {
	lo maxHeap // smaller half, root = its largest
	hi IntHeap // larger half, root = its smallest
}

func (m *MedianFinder) AddNum(x int) {
	heap.Push(&m.lo, x)
	heap.Push(&m.hi, heap.Pop(&m.lo).(int)) // pass the largest of lo across so halves stay ordered
	if m.hi.Len() > m.lo.Len() {
		heap.Push(&m.lo, heap.Pop(&m.hi).(int)) // rebalance sizes
	}
}

func (m *MedianFinder) FindMedian() float64 {
	if m.lo.Len() == 0 {
		return 0
	}
	if m.lo.Len() > m.hi.Len() {
		return float64(m.lo[0])
	}
	return float64(m.lo[0]+m.hi[0]) / 2
}