proposal: container/heap: efficient in-order Iter function

[ClickerMonkey]

Proposal Details

The recommended way to iterate over a heap in order is to copy it and pop it until it's empty. This may not be possible for some heap.Interface implementations or may be very costly (the data copying) depending on the elements in the heap. I'm proposing the following func be added to the container/heap module. It operates in O(n log n), only allocates half of the size of the heap using only ints, and only relies on Len() and Less(i,j) on the heap.Interface.

Recommended usage:

for heapIndex := range heap.Iter(myHeap) {
  // heapIndex = the index in the heap's backing structure (typically a slice)
}

A more concrete implementation using the myHeap type in heap_test.go.

// Values returns an iterator over the values in this heap implementation.
func (h *myHeap) Values() iter.Seq[int] {
	return func(yield func(int) bool) {
		for i := range heap.Iter(h) {
			if !yield((*h)[i]) {
				return
			}
		}
	}
}
// usage
for v := range myHeap.Values() {
  // the value in the heap the is not less than the one before it
}

Implementation:

// Iterates a heap interface in-order without copying.
// This is a modified breadth first search that looks at the
// children of the lowest value in the heap and yields them
// in order of their value. This will allocate an []int slice
// half the size of the heap. The underlying heap must not be
// modified during iteration.
func Iter(h Interface) iter.Seq[int] {
	n := h.Len()
	if n == 0 {
		return func(yield func(int) bool) {}
	}

	// front is all nodes in modified breadth first search
	// the breadth in the context of a heap are not items in the same
	// depth but are items all being looked at for the lowest value.
	// the lowest value is found, removed, yielded, and its children are added to the front.
	front := make([]int, (n+1)/2)
	frontSize := 1

	// adds to the front the child nodes of i
	split := func(heapIndex int) {
		if left := heapIndex*2 + 1; left < n {
			front[frontSize] = left
			frontSize++
		}
		if right := heapIndex*2 + 2; right < n {
			front[frontSize] = right
			frontSize++
		}
	}

	// takes the item at the given index in the front
	take := func(frontIndex int) int {
		taken := front[frontIndex]
		frontSize--
		front[frontIndex] = front[frontSize]
		return taken
	}

	// iterate until we've looked at all items on the heap
	return func(yield func(int) bool) {
		for frontSize > 0 {
			chosenIndex := 0
			for i := 1; i < frontSize; i++ {
				if h.Less(front[i], front[chosenIndex]) {
					chosenIndex = i
				}
			}
			chosen := take(chosenIndex)
			if !yield(chosen) {
				return
			}
			split(chosen)
		}
	}
}

First time filing a ticket and I can provide an MR with unit tests, I don't even know if these sort of suggestions are supported. Just straight up giving code. If this having an in-order implementation and other supported types not having it feels odd I can make a more comprehensive MR with whatever data structures make this feel like a more complete contribution. Otherwise I'll just keep this code to myself and anyone else interested can use it as well. I know I couldn't find any examples of code like this online and it saddened me.

[gabyhelp]

Related Issues

• proposal: container/heap: add Heap, a heap backed by a slice #47632

_{(Emoji vote if this was helpful or unhelpful; more detailed feedback welcome in this discussion.)}

[seankhliao]

Heaps generally aren't sorted data structures, providing an iterator seems wrong.
See also #47632