diff --git a/bitset.go b/bitset.go index 8fb9e9f..9aec0f3 100644 --- a/bitset.go +++ b/bitset.go @@ -94,17 +94,17 @@ func (b *BitSet) SetBitsetFrom(buf []uint64) { b.set = buf } -// From is a constructor used to create a BitSet from an array of integers +// From is a constructor used to create a BitSet from an array of words func From(buf []uint64) *BitSet { return FromWithLength(uint(len(buf))*64, buf) } -// FromWithLength constructs from an array of integers and length. +// FromWithLength constructs from an array of words and length. func FromWithLength(len uint, set []uint64) *BitSet { return &BitSet{len, set} } -// Bytes returns the bitset as array of integers +// Bytes returns the bitset as array of words func (b *BitSet) Bytes() []uint64 { return b.set } @@ -1135,3 +1135,37 @@ func (b *BitSet) UnmarshalJSON(data []byte) error { _, err = b.ReadFrom(bytes.NewReader(buf)) return err } + +// Rank returns the nunber of set bits up to and including the index +// that are set in the bitset. +// See https://en.wikipedia.org/wiki/Ranking#Ranking_in_statistics +func (b *BitSet) Rank(index uint) uint { + if index >= b.length { + return b.Count() + } + leftover := (index + 1) & 63 + answer := uint(popcntSlice(b.set[:(index+1)>>6])) + if leftover != 0 { + answer += uint(popcount(b.set[(index+1)>>6] << (64 - leftover))) + } + return answer +} + +// Select returns the index of the jth set bit, where j is the argument. +// The caller is responsible to ensure that 0 <= j < Count(): when j is +// out of range, the function returns the length of the bitset (b.length). +// +// Note that this function differs in convention from the Rank function which +// returns 1 when ranking the smallest value. We follow the conventional +// textbook definition of Select and Rank. +func (b *BitSet) Select(index uint) uint { + leftover := index + for idx, word := range b.set { + w := uint(popcount(word)) + if w > leftover { + return uint(idx)*64 + select64(word, leftover) + } + leftover -= w + } + return b.length +} diff --git a/bitset_test.go b/bitset_test.go index 3dae32e..8f50df5 100644 --- a/bitset_test.go +++ b/bitset_test.go @@ -1497,6 +1497,43 @@ func TestClearAll(t *testing.T) { } } +func TestRankSelect(t *testing.T) { + u := []uint{2, 3, 5, 7, 11, 700, 1500} + b := BitSet{} + for _, v := range u { + b.Set(v) + } + + if b.Rank(5) != 3 { + t.Error("Unexpected rank") + return + } + if b.Rank(6) != 3 { + t.Error("Unexpected rank") + return + } + if b.Rank(1500) != 7 { + t.Error("Unexpected rank") + return + } + if b.Select(0) != 2 { + t.Error("Unexpected select") + return + } + if b.Select(1) != 3 { + t.Error("Unexpected select") + return + } + if b.Select(2) != 5 { + t.Error("Unexpected select") + return + } + + if b.Select(5) != 700 { + t.Error("Unexpected select") + return + } +} func TestFlip(t *testing.T) { b := new(BitSet) c := b.Flip(11) diff --git a/select.go b/select.go new file mode 100644 index 0000000..f15e74a --- /dev/null +++ b/select.go @@ -0,0 +1,45 @@ +package bitset + +func select64(w uint64, j uint) uint { + seen := 0 + // Divide 64bit + part := w & 0xFFFFFFFF + n := uint(popcount(part)) + if n <= j { + part = w >> 32 + seen += 32 + j -= n + } + ww := part + + // Divide 32bit + part = ww & 0xFFFF + + n = uint(popcount(part)) + if n <= j { + part = ww >> 16 + seen += 16 + j -= n + } + ww = part + + // Divide 16bit + part = ww & 0xFF + n = uint(popcount(part)) + if n <= j { + part = ww >> 8 + seen += 8 + j -= n + } + ww = part + + // Lookup in final byte + counter := 0 + for ; counter < 8; counter++ { + j -= uint((ww >> counter) & 1) + if j+1 == 0 { + break + } + } + return uint(seen + counter) +}