Skip to main content

bit_manipulation/
lib.rs

1//! Bit manipulation utilities and a tiny fixed-size bitset.
2//!
3//! The free-standing functions cover the operations that show up over and
4//! over in competitive programming:
5//!
6//! * [`popcount`] — number of set bits
7//! * [`trailing_zeros`] / [`trailing_ones`] — index of the lowest set / unset bit
8//! * [`highest_bit`] / [`lowest_bit`] — read the most / least significant set bit
9//! * [`is_power_of_two`]
10//! * [`next_power_of_two`] / [`log2_floor`]
11//! * [`subsets`] — iterate over all `2^k` non-empty subsets of a `k`-bit mask
12//!
13//! [`Bitset`] wraps `Vec<u64>` and exposes the operations you'd expect:
14//! `get`, `set`, `reset`, `flip`, plus `count_ones` over the whole set.
15//!
16//! All inputs are `u64` unless documented otherwise.
17
18/// Number of set bits in `x`.
19pub fn popcount(x: u64) -> u32 {
20    x.count_ones()
21}
22
23/// Index of the lowest set bit. Returns 64 if `x == 0`.
24pub fn trailing_zeros(x: u64) -> u32 {
25    x.trailing_zeros()
26}
27
28/// Index of the lowest unset bit of `!x`. Returns 64 if `x` is all ones.
29pub fn trailing_ones(x: u64) -> u32 {
30    (!x).trailing_zeros()
31}
32
33/// Mask with the highest set bit. Returns 0 if `x == 0`.
34pub fn highest_bit(x: u64) -> u64 {
35    if x == 0 {
36        0
37    } else {
38        1u64 << (63 - x.leading_zeros())
39    }
40}
41
42/// Mask with the lowest set bit. Returns 0 if `x == 0`.
43pub fn lowest_bit(x: u64) -> u64 {
44    x & x.wrapping_neg()
45}
46
47/// True iff `x` is a (positive) power of two.
48pub fn is_power_of_two(x: u64) -> bool {
49    x != 0 && (x & (x - 1)) == 0
50}
51
52/// Smallest power of two ≥ `x`. Returns 1 for `x == 0`.
53pub fn next_power_of_two(x: u64) -> u64 {
54    if x <= 1 {
55        return 1;
56    }
57    1u64 << (64 - (x - 1).leading_zeros())
58}
59
60/// Floor of the base-2 logarithm of `x`. Panics if `x == 0`.
61pub fn log2_floor(x: u64) -> u32 {
62    assert!(x > 0, "log2_floor requires x > 0");
63    63 - x.leading_zeros()
64}
65
66/// Iterate over all `2^k` non-empty subsets of a `k`-bit mask.
67///
68/// `k` is the bit-width of the universe (e.g. 4 for a 4-bit mask).
69pub fn subsets(k: u32) -> impl Iterator<Item = u64> {
70    (1u64..(1u64 << k)).map(move |m| m)
71}
72
73/// Fix-sized bitset backed by `Vec<u64>`.
74#[derive(Debug, Clone)]
75pub struct Bitset {
76    /// Number of bits (must be ≤ `64 * words.len()`).
77    n: usize,
78    /// Storage, little-endian: bit `i` is in `words[i / 64]` at bit `i % 64`.
79    words: Vec<u64>,
80}
81
82impl Bitset {
83    /// Construct a new bitset of `n` bits, all zero.
84    pub fn new(n: usize) -> Self {
85        let nwords = n.div_ceil(64);
86        Self { n, words: vec![0; nwords] }
87    }
88
89    /// Get bit at index `i`.
90    pub fn get(&self, i: usize) -> bool {
91        assert!(i < self.n, "index out of range");
92        (self.words[i / 64] >> (i % 64)) & 1 != 0
93    }
94
95    /// Set bit at index `i` to `1`.
96    pub fn set(&mut self, i: usize) {
97        assert!(i < self.n, "index out of range");
98        self.words[i / 64] |= 1u64 << (i % 64);
99    }
100
101    /// Reset bit at index `i` to `0`.
102    pub fn reset(&mut self, i: usize) {
103        assert!(i < self.n, "index out of range");
104        self.words[i / 64] &= !(1u64 << (i % 64));
105    }
106
107    /// Flip bit at index `i`.
108    pub fn flip(&mut self, i: usize) {
109        assert!(i < self.n, "index out of range");
110        self.words[i / 64] ^= 1u64 << (i % 64);
111    }
112
113    /// Total number of set bits.
114    pub fn count_ones(&self) -> u32 {
115        self.words.iter().map(|w| w.count_ones()).sum()
116    }
117}