minheap.go

// Package heapx implements a generic binary min-heap (priority queue).
package heapx

// Ordered constrains types that support the < operator.
type Ordered interface {
	~int | ~int8 | ~int16 | ~int32 | ~int64 |
		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |
		~float32 | ~float64 | ~string
}

// MinHeap is an array-backed binary min-heap.
//
// The zero value is a usable empty heap.
type MinHeap[T Ordered] struct {
	a []T
}

// New returns a MinHeap with the given initial capacity.
func New[T Ordered](initialCap int) *MinHeap[T] {
	if initialCap < 0 {
		initialCap = 0
	}
	return &MinHeap[T]{a: make([]T, 0, initialCap)}
}

// FromSlice builds a min-heap in O(n) by Floyd's heapify on a copy of s.
func FromSlice[T Ordered](s []T) *MinHeap[T] {
	h := &MinHeap[T]{a: append([]T(nil), s...)}
	h.heapify()
	return h
}

// Len returns the number of elements in the heap.
func (h *MinHeap[T]) Len() int { return len(h.a) }

// Push inserts v into the heap in O(log n).
func (h *MinHeap[T]) Push(v T) {
	h.a = append(h.a, v)
	h.siftUp(len(h.a) - 1)
}

// Pop removes and returns the minimum element along with true. If the heap
// is empty it returns the zero value of T and false.
func (h *MinHeap[T]) Pop() (T, bool) {
	var zero T
	n := len(h.a)
	if n == 0 {
		return zero, false
	}
	min := h.a[0]
	last := n - 1
	h.a[0] = h.a[last]
	h.a[last] = zero // help GC if T is pointer-like
	h.a = h.a[:last]
	if len(h.a) > 0 {
		h.siftDown(0)
	}
	return min, true
}

// Peek returns the minimum element without removing it.
func (h *MinHeap[T]) Peek() (T, bool) {
	var zero T
	if len(h.a) == 0 {
		return zero, false
	}
	return h.a[0], true
}

// Snapshot returns a copy of the underlying array (not necessarily sorted).
func (h *MinHeap[T]) Snapshot() []T {
	out := make([]T, len(h.a))
	copy(out, h.a)
	return out
}

func (h *MinHeap[T]) siftUp(i int) {
	for i > 0 {
		p := (i - 1) / 2
		if !(h.a[i] < h.a[p]) {
			return
		}
		h.a[i], h.a[p] = h.a[p], h.a[i]
		i = p
	}
}

func (h *MinHeap[T]) siftDown(i int) {
	n := len(h.a)
	for {
		left := i*2 + 1
		if left >= n {
			return
		}
		smallest := left
		if right := left + 1; right < n && h.a[right] < h.a[left] {
			smallest = right
		}
		if !(h.a[smallest] < h.a[i]) {
			return
		}
		h.a[i], h.a[smallest] = h.a[smallest], h.a[i]
		i = smallest
	}
}

// heapify turns h.a into a valid min-heap in O(n).
func (h *MinHeap[T]) heapify() {
	for i := len(h.a)/2 - 1; i >= 0; i-- {
		h.siftDown(i)
	}
}

// IsMinHeap reports whether s satisfies the min-heap property.
func IsMinHeap[T Ordered](s []T) bool {
	for i := 0; i < len(s); i++ {
		l := i*2 + 1
		r := l + 1
		if l < len(s) && s[l] < s[i] {
			return false
		}
		if r < len(s) && s[r] < s[i] {
			return false
		}
	}
	return true
}

// KSmallest returns the k smallest elements of s in ascending order.
// Time: O(n log k). Space: O(k).
//
// The implementation uses a min-heap of size n, popping k times. For very
// large n with small k, a max-heap of size k is more memory-efficient; that
// variant is left as an exercise.
func KSmallest[T Ordered](s []T, k int) []T {
	if k <= 0 || len(s) == 0 {
		return nil
	}
	if k > len(s) {
		k = len(s)
	}
	h := FromSlice(s)
	out := make([]T, 0, k)
	for i := 0; i < k; i++ {
		v, _ := h.Pop()
		out = append(out, v)
	}
	return out
}