diff --git a/bitarray/bitarray.go b/bitarray/bitarray.go index cee0361..d2e7ec9 100644 --- a/bitarray/bitarray.go +++ b/bitarray/bitarray.go @@ -202,6 +202,15 @@ func (ba *bitArray) ClearBit(k uint64) error { return nil } +// Count returns the number of set bits in this array. +func (ba *bitArray) Count() uint64 { + count := 0 + for _, block := range ba.blocks { + count += bits.OnesCount64(uint64(block)) + } + return uint64(count) +} + // Or will bitwise or two bit arrays and return a new bit array // representing the result. func (ba *bitArray) Or(other BitArray) BitArray { diff --git a/bitarray/sparse_bitarray.go b/bitarray/sparse_bitarray.go index 52653b0..19ec140 100644 --- a/bitarray/sparse_bitarray.go +++ b/bitarray/sparse_bitarray.go @@ -16,7 +16,10 @@ limitations under the License. package bitarray -import "sort" +import ( + "math/bits" + "sort" +) // uintSlice is an alias for a slice of ints. Len, Swap, and Less // are exported to fulfill an interface needed for the search @@ -243,6 +246,15 @@ func (sba *sparseBitArray) Equals(other BitArray) bool { return true } +// Count returns the number of set bits in this array. +func (sba *sparseBitArray) Count() uint64 { + count := 0 + for _, block := range sba.blocks { + count += bits.OnesCount64(uint64(block)) + } + return uint64(count) +} + // Or will perform a bitwise or operation with the provided bitarray and // return a new result bitarray. func (sba *sparseBitArray) Or(other BitArray) BitArray {