implement a Ceil func

treap.go

@@ -83,6 +83,31 @@ func (t *Tree) Len() int {
 	return t.count
 }
 
+// Get the first key, item >= key in the tree.
+func (t *Tree) Ceil(key Key) (k Key, v Item) {
+	t.ceil(t.root, key, &k, &v, false)
+	return k, v
+}
+
+func (t *Tree) ceil(node *Node, key Key, bestK *Key, bestV *Item, found bool) {
+	if node == nil {
+		return
+	}
+	// This found var is neccessary because the empty value of key varies (nil, 0, ""),
+	// so I can't test for it
+	if !t.less(node.key, key) && (!found || t.less(node.key, *bestK)) {
+		found = true
+		*bestK = node.key
+		*bestV = node.item
+	}
+	if t.less(key, node.key) {
+		t.ceil(node.left, key, bestK, bestV, found)
+	}
+	if t.less(node.key, key) {
+		t.ceil(node.right, key, bestK, bestV, found)
+	}
+}
+
 // Get an Item in the tree.
 func (t *Tree) Get(key Key) Item {
 	return t.get(t.root, key)

treap_test.go

@@ -40,6 +40,45 @@ func TestInsert(t *testing.T) {
 	}
 }
 
+func TestCeil(t *testing.T) {
+	tree := NewTree(StringLess)
+	tree.Insert("xyz", "adsf")
+	_, y := tree.Ceil("x")
+	if y != "adsf" {
+		t.Errorf("expected adsf, got %v", y)
+	}
+
+	x, _ := tree.Ceil("y")
+	if x != nil {
+		t.Errorf("expected nil, got %v", x)
+	}
+}
+
+func TestCeilPermutations(t *testing.T) {
+	rounds := 100
+	for r := 0; r < rounds; r++ {
+		n := 30
+		tree := NewTree(IntLess)
+		ints := rand.Perm(n)
+		for _, i := range ints {
+			tree.Insert(i*2, i*2)
+		}
+
+		k, _ := tree.Ceil(n * 2)
+		if k != nil {
+			t.Errorf("expected nil, got %d", k)
+		}
+
+		for i := 0; i < n-1; i++ {
+			j := i*2 + 1
+			k, _ := tree.Ceil(j)
+			if k != j+1 {
+				t.Errorf("expected %d, got %d", j+1, k)
+			}
+		}
+	}
+}
+
 func TestFromDoc(t *testing.T) {
 	tree := NewTree(IntLess)
 	tree.Insert(5, "a")