diff --git a/container/intsets/sparse.go b/container/intsets/sparse.go index abe172d34f..5db01c1a44 100644 --- a/container/intsets/sparse.go +++ b/container/intsets/sparse.go @@ -170,6 +170,26 @@ func (b *block) min(take bool) int { 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 @@ func (s *Sparse) Min() int { 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 d9d4036aee..e3ef9d3d05 100644 --- a/container/intsets/sparse_test.go +++ b/container/intsets/sparse_test.go @@ -275,11 +275,30 @@ func TestRandomMutations(t *testing.T) { } } +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<= 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.