Tournament Tree aka Loser Tree, in Go

tree.go

// Loser tree, from https://en.wikipedia.org/wiki/K-way_merge_algorithm#Tournament_Tree

package loser

import (
	"math"
)

// Ideally Ordered would be any int, float, etc., bit we need to have a maxVal too, so kludge it.
type Ordered interface{ ~uint64 }

const maxVal = uint64(math.MaxUint64) // This value must be higher than all real values used in the tree.

type Sequence[E Ordered] interface {
	Next() bool    // Advances and returns true if there is a value at this new position.
	Seek(v E) bool // Advances the iterator to value v or greater and returns true if a value was found.
	At() E         // Returns the value at the current position.
}

func NewMerge[E Ordered](nElements int) *LoserTree[E] {
	t := LoserTree[E]{
		nodes: make([]node[E], nElements, nElements*2),
	}
	if nElements > 0 {
		t.nodes[0].index = -1 // flag to be initialized on first call to Next().
	}
	return &t
}

// Must be called exactly nElements times after calling NewMerge.
func (t *LoserTree[E]) Add(e Sequence[E]) {
	t.nodes = append(t.nodes, node[E]{items: e})
}

// A loser tree is a binary tree laid out such that nodes N and N+1 have parent N/2.
// We store M leaf nodes in positions M...2M-1, and M-1 internal nodes in positions 1..M-1.
// Node 0 is a special node, containing the winner of the contest.
type LoserTree[E Ordered] struct {
	nodes []node[E]
}

type node[E Ordered] struct {
	index int         // This is the loser for all nodes except the 0th, where it is the winner.
	value E           // Value copied from the loser node, or winner for node 0.
	items Sequence[E] // Only populated for leaf nodes.
}

func (n *node[E]) moveNext() bool {
	ret := n.items.Next()
	if ret {
		n.value = n.items.At()
	} else {
		n.value = E(maxVal)
	}
	return ret
}

func (t *LoserTree[E]) At() E {
	return t.nodes[0].value
}

func (t *LoserTree[E]) Next() bool {
	if len(t.nodes) == 0 {
		return false
	}
	if t.nodes[0].index == -1 { // If tree has not been initialized yet, do that.
		t.initialize()
		return t.nodes[0].value != E(maxVal)
	}
	prev := t.nodes[0].value
	for {
		t.nodes[t.nodes[0].index].moveNext()
		t.replayGames(t.nodes[0].index)
		if t.nodes[0].value != prev {
			break
		}
	}
	return t.nodes[0].value != E(maxVal)
}

func (t *LoserTree[E]) Seek(val E) bool {
	if len(t.nodes) == 0 {
		return false
	}
	if t.nodes[0].index == -1 { // If tree has not been initialized yet, do that.
		t.initialize()
	}
	// While the current winner is before the target val, seek that one, then find the new lowest value.
	for t.nodes[0].value < val {
		cur := &t.nodes[t.nodes[0].index]
		if !cur.items.Seek(val) {
			cur.value = E(maxVal)
		} else {
			cur.value = cur.items.At()
		}
		t.replayGames(t.nodes[0].index)
	}
	return t.nodes[0].value != E(maxVal)
}

func (t *LoserTree[E]) initialize() {
	winners := make([]int, len(t.nodes))
	// Initialize leaf nodes as winners to start.
	for i := len(t.nodes) / 2; i < len(t.nodes); i++ {
		winners[i] = i
		t.nodes[i].moveNext() // Must call Next on each item so that At() has a value.
	}
	for i := len(t.nodes) - 2; i > 0; i -= 2 {
		// At each stage the winners play each other, and we record the loser in the node.
		loser, winner := t.playGame(winners[i], winners[i+1])
		p := parent(i)
		t.nodes[p].index = loser
		t.nodes[p].value = t.nodes[loser].value
		winners[p] = winner
	}
	t.nodes[0].index = winners[1]
	t.nodes[0].value = t.nodes[winners[1]].value
}

// Starting at pos, re-consider all values up to the root.
func (t *LoserTree[E]) replayGames(pos int) {
	// At the start, pos is a leaf node, and is the winner at that level.
	n := parent(pos)
	for n != 0 {
		if t.nodes[n].value < t.nodes[pos].value {
			loser := pos
			// Record pos as the loser here, and the old loser is the new winner.
			pos = t.nodes[n].index
			t.nodes[n].index = loser
			t.nodes[n].value = t.nodes[loser].value
		}
		n = parent(n)
	}
	// pos is now the winner; store it in node 0.
	t.nodes[0].index = pos
	t.nodes[0].value = t.nodes[pos].value
}

func (t *LoserTree[E]) playGame(a, b int) (loser, winner int) {
	if t.nodes[a].value < t.nodes[b].value {
		return b, a
	}
	return a, b
}

func parent(i int) int { return i / 2 }