Merge pull request #79 from contain-rs/hamming-ones-zeros

Hamming distance: count_ones and count_zeros

src/lib.rs

@@ -142,6 +142,10 @@ pub trait BitBlock:
     fn from_byte(byte: u8) -> Self;
     /// Count the number of 1's in the bitwise repr
     fn count_ones(self) -> usize;
+    /// Count the number of 0's in the bitwise repr
+    fn count_zeros(self) -> usize {
+        Self::bits() - self.count_ones()
+    }
     /// Get `0`
     fn zero() -> Self;
     /// Get `1`
@@ -158,6 +162,8 @@ macro_rules! bit_block_impl {
             #[inline]
             fn count_ones(self) -> usize { self.count_ones() as usize }
             #[inline]
+            fn count_zeros(self) -> usize { self.count_zeros() as usize }
+            #[inline]
             fn one() -> Self { 1 }
             #[inline]
             fn zero() -> Self { 0 }
@@ -401,7 +407,7 @@ impl<B: BitBlock> BitVec<B> {
         changed_bits != B::zero()
     }
 
-    /// Iterator over mutable refs to  the underlying blocks of data.
+    /// Iterator over mutable refs to the underlying blocks of data.
     #[inline]
     fn blocks_mut(&mut self) -> MutBlocks<B> {
         // (2)
@@ -931,6 +937,56 @@ impl<B: BitBlock> BitVec<B> {
         }) && (last_word == mask_for_bits(self.nbits))
     }
 
+    /// Returns the number of ones in the binary representation.
+    ///
+    /// Also known as the
+    /// [Hamming weight](https://en.wikipedia.org/wiki/Hamming_weight).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bit_vec::BitVec;
+    ///
+    /// let mut bv = BitVec::from_elem(100, true);
+    /// assert_eq!(bv.count_ones(), 100);
+    ///
+    /// bv.set(50, false);
+    /// assert_eq!(bv.count_ones(), 99);
+    /// ```
+    #[inline]
+    pub fn count_ones(&self) -> u64 {
+        self.ensure_invariant();
+        // Add the number of ones of each block.
+        self.blocks().map(|elem| elem.count_ones() as u64).sum()
+    }
+
+    /// Returns the number of zeros in the binary representation.
+    ///
+    /// Also known as the opposite of
+    /// [Hamming weight](https://en.wikipedia.org/wiki/Hamming_weight).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bit_vec::BitVec;
+    ///
+    /// let mut bv = BitVec::from_elem(100, false);
+    /// assert_eq!(bv.count_zeros(), 100);
+    ///
+    /// bv.set(50, true);
+    /// assert_eq!(bv.count_zeros(), 99);
+    /// ```
+    #[inline]
+    pub fn count_zeros(&self) -> u64 {
+        self.ensure_invariant();
+        // Add the number of zeros of each block.
+        let extra_zeros = (B::bits() - (self.len() % B::bits())) % B::bits();
+        self.blocks()
+            .map(|elem| elem.count_zeros() as u64)
+            .sum::<u64>()
+            - extra_zeros as u64
+    }
+
     /// Returns an iterator over the elements of the vector in order.
     ///
     /// # Examples
@@ -2637,4 +2693,40 @@ mod tests {
         a.append(&mut c);
         assert_eq!(&[0xc0, 0x00, 0x00, 0x00, 0x3f, 0xc0][..], &*a.to_bytes());
     }
+
+    #[test]
+    fn test_count_ones() {
+        for i in 0..1000 {
+            let mut t = BitVec::from_elem(i, true);
+            let mut f = BitVec::from_elem(i, false);
+            assert_eq!(i as u64, t.count_ones());
+            assert_eq!(0 as u64, f.count_ones());
+            if i > 20 {
+                t.set(10, false);
+                t.set(i - 10, false);
+                assert_eq!(i - 2, t.count_ones() as usize);
+                f.set(10, true);
+                f.set(i - 10, true);
+                assert_eq!(2, f.count_ones());
+            }
+        }
+    }
+
+    #[test]
+    fn test_count_zeros() {
+        for i in 0..1000 {
+            let mut tbits = BitVec::from_elem(i, true);
+            let mut fbits = BitVec::from_elem(i, false);
+            assert_eq!(i as u64, fbits.count_zeros());
+            assert_eq!(0 as u64, tbits.count_zeros());
+            if i > 20 {
+                fbits.set(10, true);
+                fbits.set(i - 10, true);
+                assert_eq!(i - 2, fbits.count_zeros() as usize);
+                tbits.set(10, false);
+                tbits.set(i - 10, false);
+                assert_eq!(2, tbits.count_zeros());
+            }
+        }
+    }
 }