snoob.rs

/// Returns the next smallest integer with the same hamming weight as the input.
pub fn snoob(x: u32) -> u32 {
    let mut ntz = x.trailing_zeros();
    if x == 0 {
        ntz = 0;
    }
    let least_one = x & x.wrapping_neg(); // nightly: x.isolate_least_significant_one()
    let ripple = x + least_one;
    let ones = (x ^ ripple).unbounded_shr(2 + ntz);
    ripple | ones
}

/// Iterates integers with the same hamming weight within a certain number of bits.
pub struct BinarySubsetIter {
    next: Option<u32>,
    end: u32,
}

impl Iterator for BinarySubsetIter {
    type Item = u32;
    fn next(&mut self) -> Option<Self::Item> {
        let next = self.next?;
        if next == self.end {
            self.next = None;
        } else {
            self.next = Some(snoob(next));
        }
        Some(next)
    }
}

/// Returns an iterator that yields, in increasing order, each integer such that:
///
///   - No bits beyond the first `n` are set.
///   - The hamming weight ([`u32::count_ones`]) equals `k`.
///
pub fn binary_subsets(n: u32, k: u32) -> BinarySubsetIter {
    assert!(n > 0);
    assert!(n <= 32);
    assert!(k <= n);
    let start = u32::MAX.unbounded_shr(32 - k);
    BinarySubsetIter {
        next: Some(start),
        end: start.unbounded_shl(n - k),
    }
}

#[test]
fn test_binary_subsets() {
    let mut iter = binary_subsets(1, 1);
    assert_eq!(iter.next(), Some(1));
    assert_eq!(iter.next(), None);

    let mut iter = binary_subsets(3, 1);
    assert_eq!(iter.next(), Some(1));
    assert_eq!(iter.next(), Some(2));
    assert_eq!(iter.next(), Some(4));
    assert_eq!(iter.next(), None);

    let mut iter = binary_subsets(8, 3);
    assert_eq!(iter.next(), Some(0b00000111));
    assert_eq!(iter.next(), Some(0b00001011));
    assert_eq!(iter.next(), Some(0b00001101));
    assert_eq!(iter.next(), Some(0b00001110));
    assert_eq!(iter.next(), Some(0b00010011));
    assert_eq!(iter.next(), Some(0b00010101));
    assert_eq!(iter.next(), Some(0b00010110));
    assert_eq!(iter.next(), Some(0b00011001));
    assert_eq!(iter.next(), Some(0b00011010));
    assert_eq!(iter.next(), Some(0b00011100));
    assert_eq!(iter.next(), Some(0b00100011));
    assert_eq!(iter.next(), Some(0b00100101));
    assert_eq!(iter.next(), Some(0b00100110));
    assert_eq!(iter.next(), Some(0b00101001));
    assert_eq!(iter.next(), Some(0b00101010));
    assert_eq!(iter.next(), Some(0b00101100));
    assert_eq!(iter.next(), Some(0b00110001));
    assert_eq!(iter.next(), Some(0b00110010));
    assert_eq!(iter.next(), Some(0b00110100));
    assert_eq!(iter.next(), Some(0b00111000));
    assert_eq!(iter.next(), Some(0b01000011));
    assert_eq!(iter.next(), Some(0b01000101));
    assert_eq!(iter.next(), Some(0b01000110));
    assert_eq!(iter.next(), Some(0b01001001));
    assert_eq!(iter.next(), Some(0b01001010));
    assert_eq!(iter.next(), Some(0b01001100));
    assert_eq!(iter.next(), Some(0b01010001));
    assert_eq!(iter.next(), Some(0b01010010));
    assert_eq!(iter.next(), Some(0b01010100));
    assert_eq!(iter.next(), Some(0b01011000));
    assert_eq!(iter.next(), Some(0b01100001));
    assert_eq!(iter.next(), Some(0b01100010));
    assert_eq!(iter.next(), Some(0b01100100));
    assert_eq!(iter.next(), Some(0b01101000));
    assert_eq!(iter.next(), Some(0b01110000));
    assert_eq!(iter.next(), Some(0b10000011));
    assert_eq!(iter.next(), Some(0b10000101));
    assert_eq!(iter.next(), Some(0b10000110));
    assert_eq!(iter.next(), Some(0b10001001));
    assert_eq!(iter.next(), Some(0b10001010));
    assert_eq!(iter.next(), Some(0b10001100));
    assert_eq!(iter.next(), Some(0b10010001));
    assert_eq!(iter.next(), Some(0b10010010));
    assert_eq!(iter.next(), Some(0b10010100));
    assert_eq!(iter.next(), Some(0b10011000));
    assert_eq!(iter.next(), Some(0b10100001));
    assert_eq!(iter.next(), Some(0b10100010));
    assert_eq!(iter.next(), Some(0b10100100));
    assert_eq!(iter.next(), Some(0b10101000));
    assert_eq!(iter.next(), Some(0b10110000));
    assert_eq!(iter.next(), Some(0b11000001));
    assert_eq!(iter.next(), Some(0b11000010));
    assert_eq!(iter.next(), Some(0b11000100));
    assert_eq!(iter.next(), Some(0b11001000));
    assert_eq!(iter.next(), Some(0b11010000));
    assert_eq!(iter.next(), Some(0b11100000));
    assert_eq!(iter.next(), None);

    for n in 1..=32 {
        let mut iter_a = binary_subsets(n, 0);
        let mut iter_b = binary_subsets(n, n);
        assert_eq!(iter_a.next(), Some(0));
        assert_eq!(iter_b.next(), Some(u32::MAX >> (32 - n)));
        assert_eq!(iter_a.next(), None);
        assert_eq!(iter_b.next(), None);
    }

    let mut iter_a = binary_subsets(32, 1);
    let mut iter_b = binary_subsets(32, 31);
    for i in 0..=31 {
        assert_eq!(iter_a.next(), Some(1 << i));
        assert_eq!(iter_b.next(), Some(u32::MAX ^ (1 << (31 - i))));
    }
    assert_eq!(iter_a.next(), None);
    assert_eq!(iter_b.next(), None);
}