Go

Implement a generic binary search tree keyed by cmp.Ordered with an in-order iterator.

Question 133MediumGo 1.22 to 1.25
type Tree[K cmp.Ordered, V any] struct {
    root *node[K, V]
    size int
}

type node[K cmp.Ordered, V any] struct {
    key         K
    val         V
    left, right *node[K, V]
}

func (t *Tree[K, V]) Put(k K, v V) {
    p := &t.root
    for *p != nil {
        switch c := cmp.Compare(k, (*p).key); {
        case c < 0:
            p = &(*p).left
        case c > 0:
            p = &(*p).right
        default:
            (*p).val = v
            return
        }
    }
    *p = &node[K, V]{key: k, val: v}
    t.size++
}

func (t *Tree[K, V]) All() iter.Seq2[K, V] {
    return func(yield func(K, V) bool) { t.root.walk(yield) }
}

func (n *node[K, V]) walk(yield func(K, V) bool) bool {
    if n == nil {
        return true
    }
    return n.left.walk(yield) && yield(n.key, n.val) && n.right.walk(yield)
}

// usage
var t Tree[string, int]
t.Put("b", 2); t.Put("a", 1); t.Put("c", 3)
for k, v := range t.All() {
    fmt.Println(k, v) // a 1, b 2, c 3
}

Talking points:

  • The pointer-to-pointer insertion avoids special-casing the root.
  • cmp.Compare orders NaN consistently. Plain < does not.
  • Returning bool from walk propagates an early break from the caller's loop.
  • For keys that are not ordered, accept a comparator: NewTree[K, V any](cmp func(K, K) int).

More on Generics

All 36 Generics questions