haar_lib/algo/enum_bit/
subset_size_k.rs

1//! ビットが`1`の個数が`k`であるものを列挙する
2use std::iter::successors;
3
4/// 幅`width`のなかで、ビットが`1`の個数が`k`であるものを列挙するイテレータを返す。
5pub fn subset_size_k(width: u32, k: u32) -> impl Iterator<Item = u32> {
6    successors(Some((1 << k) - 1), move |&t| {
7        let x = ((t as i32) & (-(t as i32))) as u32;
8        let y = t + x;
9        let t = (((t & !y) / x) >> 1) | y;
10        (t < 1 << width).then_some(t)
11    })
12}
13
14#[cfg(test)]
15mod tests {
16    use super::*;
17
18    fn check(n: u32, k: u32) {
19        let a = (0_u32..1 << n)
20            .filter(|&i| i.count_ones() == k)
21            .collect::<Vec<_>>();
22
23        let b = subset_size_k(n, k).collect::<Vec<_>>();
24
25        assert_eq!(a, b);
26    }
27
28    #[test]
29    fn test() {
30        check(10, 3);
31    }
32}