use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashSet};
pub struct BreadthFirstCombinations<const K: usize> {
items: Vec<usize>,
n: usize,
heap: BinaryHeap<Reverse<(usize, [u32; K])>>,
seen: HashSet<[u32; K]>,
}
impl<const K: usize> BreadthFirstCombinations<K> {
pub fn new(items: &[usize]) -> Self {
let n = items.len();
let mut bfc = Self {
items: items.to_vec(),
n,
heap: BinaryHeap::new(),
seen: HashSet::new(),
};
if n >= K && K > 0 {
let initial: [u32; K] = std::array::from_fn(|i| i as u32);
let sum: usize = initial.iter().map(|&x| x as usize).sum();
bfc.seen.insert(initial);
bfc.heap.push(Reverse((sum, initial)));
}
bfc
}
}
impl<const K: usize> Iterator for BreadthFirstCombinations<K> {
type Item = Vec<usize>;
fn next(&mut self) -> Option<Vec<usize>> {
let Reverse((_, combo)) = self.heap.pop()?;
for i in 0..K {
let next_val = combo[i] + 1;
let upper = if i + 1 < K {
combo[i + 1]
} else {
self.n as u32
};
if next_val < upper {
let mut new_combo = combo;
new_combo[i] = next_val;
if self.seen.insert(new_combo) {
let sum: usize = new_combo.iter().map(|&x| x as usize).sum();
self.heap.push(Reverse((sum, new_combo)));
}
}
}
Some(combo.iter().map(|&i| self.items[i as usize]).collect())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bfs_combinations_complete() {
let items: Vec<usize> = vec![10, 20, 30, 40, 50];
let combos: Vec<Vec<usize>> = BreadthFirstCombinations::<3>::new(&items).collect();
assert_eq!(combos.len(), 10);
assert_eq!(combos[0], vec![10, 20, 30]);
}
#[test]
fn test_bfs_combinations_order() {
let items: Vec<usize> = (0..6).collect();
let combos: Vec<Vec<usize>> = BreadthFirstCombinations::<4>::new(&items).collect();
assert_eq!(combos.len(), 15);
let sums: Vec<usize> = combos.iter().map(|c| c.iter().sum()).collect();
for w in sums.windows(2) {
assert!(w[0] <= w[1], "sums not in order: {} > {}", w[0], w[1]);
}
}
#[test]
fn test_bfs_combinations_early_stop() {
let items: Vec<usize> = (0..50).collect();
let first_10: Vec<Vec<usize>> =
BreadthFirstCombinations::<4>::new(&items).take(10).collect();
assert_eq!(first_10.len(), 10);
assert_eq!(first_10[0], vec![0, 1, 2, 3]);
}
#[test]
fn test_bfs_too_few_items() {
let items: Vec<usize> = vec![1, 2];
let combos: Vec<Vec<usize>> = BreadthFirstCombinations::<4>::new(&items).collect();
assert!(combos.is_empty());
}
}