diff --git a/LICENSE b/LICENSE index f65a8ec..61b445a 100644 --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,5 @@ Copyright (C) 2011 Numerotron Inc. +Copyright (C) 2013 Onemorecloud Inc. Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/treap.go b/treap.go index 2b54106..5da47f7 100644 --- a/treap.go +++ b/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) diff --git a/treap_test.go b/treap_test.go index e843877..54f9117 100644 --- a/treap_test.go +++ b/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")