diff --git a/src/permutation.rs b/src/permutation.rs index fe7d4cd..4594ca5 100644 --- a/src/permutation.rs +++ b/src/permutation.rs @@ -42,42 +42,82 @@ pub fn nth_lex(mut digits: Vec, nth: usize) -> Result, #[derive(Clone)] pub struct Permutator { - pub current: Vec, + pub elements: Vec, idx: usize, } impl Permutator { pub fn new(elements: Vec) -> Self { - Self { - current: elements, - idx: 0, - } + Self { elements, idx: 0 } } } -impl Iterator for Permutator { +impl Iterator for Permutator { type Item = Vec; /// Returns the next permutation and changes the current permutation to it /// This operation wraps around fn next(&mut self) -> Option { - if self.current.is_empty() { + if self.elements.is_empty() { return None; } - let mut digits = self.current.clone(); - if self.idx == factorial(digits.len()) { + if self.idx == 0 { + self.idx += 1; + return Some(self.elements.clone()); + } + let num_unique_digits = self.elements.len(); + if self.idx == factorial(num_unique_digits) { return None; } + let mut digits = self.elements.clone(); let mut perm = Vec::new(); - let num_unique_digits = digits.len(); - let mut remainder = 1; - for idx in 1..=digits.len() { + let mut remainder = self.idx; + for idx in 1..=num_unique_digits { let permutations = remainder / factorial(num_unique_digits - idx); remainder %= factorial(num_unique_digits - idx); - perm.push(digits[permutations].clone()); - digits.remove(permutations); + perm.push(digits.remove(permutations)); } self.idx += 1; - self.current = digits; - Some(self.current.clone()) + Some(perm) + } +} + +#[cfg(test)] +mod test { + use crate::permutation::{nth_lex, Permutator}; + + const SMALL: [i8; 3] = [1, 2, 3]; + + // Nth Lex + + #[test] + fn nth_lex_zero_small() { + let small = SMALL.to_vec(); + assert!(nth_lex(small, 0).is_err()); + } + + #[test] + fn nth_lex_small() { + let small = SMALL.to_vec(); + assert_eq!(nth_lex(small.clone(), 1).unwrap(), vec![1, 2, 3]); + assert_eq!(nth_lex(small.clone(), 2).unwrap(), vec![1, 3, 2]); + assert_eq!(nth_lex(small.clone(), 3).unwrap(), vec![2, 1, 3]); + assert_eq!(nth_lex(small.clone(), 4).unwrap(), vec![2, 3, 1]); + assert_eq!(nth_lex(small.clone(), 5).unwrap(), vec![3, 1, 2]); + assert_eq!(nth_lex(small.clone(), 6).unwrap(), vec![3, 2, 1]); + assert!(nth_lex(small.clone(), 7).is_err()); + } + + // Permutator + #[test] + fn perm_small() { + let small = SMALL.to_vec(); + let mut comb = Permutator::new(small); + assert_eq!(comb.next(), Some(vec![1, 2, 3])); + assert_eq!(comb.next(), Some(vec![1, 3, 2])); + assert_eq!(comb.next(), Some(vec![2, 1, 3])); + assert_eq!(comb.next(), Some(vec![2, 3, 1])); + assert_eq!(comb.next(), Some(vec![3, 1, 2])); + assert_eq!(comb.next(), Some(vec![3, 2, 1])); + assert_eq!(comb.next(), None); } }