x/tools/container/intsets: add LowerBound
Fixes golang/go#21310.
Change-Id: Id3f23a66b9889a5087c1f83e7d672d14c41a59e3
Reviewed-on: https://go-review.googlesource.com/53432
Reviewed-by: Alan Donovan <adonovan@google.com>
diff --git a/container/intsets/sparse.go b/container/intsets/sparse.go
index abe172d..5db01c1 100644
--- a/container/intsets/sparse.go
+++ b/container/intsets/sparse.go
@@ -170,6 +170,26 @@
panic("BUG: empty block")
}
+// lowerBound returns the smallest element of the block that is greater than or
+// equal to the element corresponding to the ith bit. If there is no such
+// element, the second return value is false.
+func (b *block) lowerBound(i uint) (int, bool) {
+ w := i / bitsPerWord
+ bit := i % bitsPerWord
+
+ if val := b.bits[w] >> bit; val != 0 {
+ return b.offset + int(i) + ntz(val), true
+ }
+
+ for w++; w < wordsPerBlock; w++ {
+ if val := b.bits[w]; val != 0 {
+ return b.offset + int(w*bitsPerWord) + ntz(val), true
+ }
+ }
+
+ return 0, false
+}
+
// forEach calls f for each element of block b.
// f must not mutate b's enclosing Sparse.
func (b *block) forEach(f func(int)) {
@@ -280,8 +300,26 @@
return s.root.min(false)
}
+// LowerBound returns the smallest element >= x, or MaxInt if there is no such
+// element.
+func (s *Sparse) LowerBound(x int) int {
+ offset, i := offsetAndBitIndex(x)
+ for b := s.first(); b != &none; b = s.next(b) {
+ if b.offset > offset {
+ return b.min(false)
+ }
+ if b.offset == offset {
+ if y, ok := b.lowerBound(i); ok {
+ return y
+ }
+ }
+ }
+ return MaxInt
+}
+
// block returns the block that would contain offset,
// or nil if s contains no such block.
+// Precondition: offset is a multiple of bitsPerBlock.
func (s *Sparse) block(offset int) *block {
for b := s.first(); b != &none && b.offset <= offset; b = s.next(b) {
if b.offset == offset {
diff --git a/container/intsets/sparse_test.go b/container/intsets/sparse_test.go
index d9d4036..e3ef9d3 100644
--- a/container/intsets/sparse_test.go
+++ b/container/intsets/sparse_test.go
@@ -275,11 +275,30 @@
}
}
+func TestLowerBound(t *testing.T) {
+ // Use random sets of sizes from 0 to about 4000.
+ prng := rand.New(rand.NewSource(0))
+ for i := uint(0); i < 12; i++ {
+ x := randomPset(prng, 1<<i)
+ for j := 0; j < 10000; j++ {
+ found := intsets.MaxInt
+ for e := range x.hash {
+ if e >= j && e < found {
+ found = e
+ }
+ }
+ if res := x.bits.LowerBound(j); res != found {
+ t.Errorf("%s: LowerBound(%d)=%d, expected %d", &x.bits, j, res, found)
+ }
+ }
+ }
+}
+
// TestSetOperations exercises classic set operations: ∩ , ∪, \.
func TestSetOperations(t *testing.T) {
prng := rand.New(rand.NewSource(0))
- // Use random sets of sizes from 0 to about 1000.
+ // Use random sets of sizes from 0 to about 4000.
// For each operator, we test variations such as
// Z.op(X, Y), Z.op(X, Z) and Z.op(Z, Y) to exercise
// the degenerate cases of each method implementation.
