Generate k-permutations in Rust

k_permutations.rs

// rustc 1.48.0 (7eac88abb 2020-11-16)

struct PermutationIterator<T> {
    v: Vec<(usize, T)>,
    k: usize,
    finished: bool,
}

impl<T> PermutationIterator<T> {
    fn new<I: Iterator<Item = T>>(iter: I, k: usize) -> PermutationIterator<T> {
        let v: Vec<(usize, T)> = iter.enumerate().collect();
        if k > v.len() {
            panic!(
                "k must be smaller than vec size. k: {}, vec size: {}",
                k,
                v.len()
            );
        }
        let finished = v.is_empty() || k == 0;
        PermutationIterator { v, k, finished }
    }

    fn flip(&mut self, mut lo: usize, mut hi: usize) {
        while lo + 1 < hi {
            self.v.swap(lo, hi - 1);
            lo += 1;
            hi -= 1;
        }
    }

    // ref: https://stackoverflow.com/questions/51287171/generating-k-permutations-from-n-in-c#answer-51292710
    fn next_k_permutation(&mut self) {
        let mut tail_max = self.v[self.n() - 1].0;
        let mut tail_start = self.k;
        while tail_start > 0 && self.v[tail_start - 1].0 >= tail_max {
            tail_start -= 1;
            tail_max = self.v[tail_start].0;
        }
        if tail_start == 0 {
            self.finished = true;
        } else {
            let mut swap_in: usize;
            let pivot = self.v[tail_start - 1].0;
            if pivot >= self.v[self.n() - 1].0 {
                swap_in = tail_start;
                while swap_in + 1 < self.k && self.v[swap_in + 1].0 > pivot {
                    swap_in += 1;
                }
            } else {
                swap_in = self.n() - 1;
                while swap_in > self.k && self.v[swap_in - 1].0 > pivot {
                    swap_in -= 1;
                }
            }
            self.v.swap(tail_start - 1, swap_in);
            self.flip(self.k, self.n());
            self.flip(tail_start, self.n());
        }
    }

    fn n(&self) -> usize {
        self.v.len()
    }
}

impl<T: Copy> Iterator for PermutationIterator<T> {
    type Item = Vec<T>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.finished {
            None
        } else {
            let res = self.v.iter().map(|e| e.1).take(self.k).collect::<Vec<T>>();
            // let res = self.v.iter().map(|e| e.1).collect::<Vec<T>>();
            self.next_k_permutation();
            Some(res)
        }
    }
}

trait Permutations<T> {
    fn permutations(self, size: usize) -> PermutationIterator<T>;
}

impl<T, I: IntoIterator<Item = T>> Permutations<T> for I {
    fn permutations(self, k: usize) -> PermutationIterator<T> {
        PermutationIterator::new(self.into_iter(), k)
    }
}

#[cfg(test)]
mod tests {
    use crate::Permutations;

    #[test]
    fn n_1_k_0() {
        let actual: Vec<Vec<i32>> = (0..1).permutations(0).collect();
        let expected: Vec<Vec<i32>> = vec![];
        assert_eq!(expected, actual);
    }

    #[test]
    fn n_1_k_1() {
        let actual: Vec<Vec<i32>> = (0..1).permutations(1).collect();
        let expected = vec![vec![0]];
        assert_eq!(expected, actual);
    }

    #[test]
    #[should_panic]
    fn n_1_k_2() {
        let _actual: Vec<Vec<i32>> = (0..1).permutations(2).collect();
    }

    #[test]
    fn n_3_k_1() {
        let actual: Vec<Vec<i32>> = (0..3).permutations(1).collect();
        let expected = vec![vec![0], vec![1], vec![2]];
        assert_eq!(expected, actual);
    }

    #[test]
    fn n_3_k_2() {
        let actual: Vec<Vec<i32>> = (0..3).permutations(2).collect();
        let expected = vec![
            vec![0, 1],
            vec![0, 2],
            vec![1, 0],
            vec![1, 2],
            vec![2, 0],
            vec![2, 1],
        ];
        assert_eq!(expected, actual);
    }

    #[test]
    fn n_3_k_3() {
        let actual: Vec<Vec<i32>> = (0..3).permutations(3).collect();
        let expected = vec![
            vec![0, 1, 2],
            vec![0, 2, 1],
            vec![1, 0, 2],
            vec![1, 2, 0],
            vec![2, 0, 1],
            vec![2, 1, 0],
        ];
        assert_eq!(expected, actual);
    }

    #[test]
    fn n_5_k_2_chars() {
        let actual: Vec<Vec<char>> = "hello".chars().permutations(2).collect();
        let expected = vec![
            vec!['h', 'e'],
            vec!['h', 'l'],
            vec!['h', 'l'],
            vec!['h', 'o'],
            vec!['e', 'h'],
            vec!['e', 'l'],
            vec!['e', 'l'],
            vec!['e', 'o'],
            vec!['l', 'h'],
            vec!['l', 'e'],
            vec!['l', 'l'],
            vec!['l', 'o'],
            vec!['l', 'h'],
            vec!['l', 'e'],
            vec!['l', 'l'],
            vec!['l', 'o'],
            vec!['o', 'h'],
            vec!['o', 'e'],
            vec!['o', 'l'],
            vec!['o', 'l'],
        ];
        assert_eq!(expected, actual);
    }

    #[test]
    fn n_3_k_2_mystruct() {
        #[derive(Debug, Copy, Clone, PartialEq)]
        struct MyStruct(i32);
        let actual: Vec<Vec<MyStruct>> = vec![MyStruct(0), MyStruct(1), MyStruct(2)]
            .permutations(2)
            .collect();
        let expected = vec![
            vec![MyStruct(0), MyStruct(1)],
            vec![MyStruct(0), MyStruct(2)],
            vec![MyStruct(1), MyStruct(0)],
            vec![MyStruct(1), MyStruct(2)],
            vec![MyStruct(2), MyStruct(0)],
            vec![MyStruct(2), MyStruct(1)],
        ];
        assert_eq!(expected, actual);
    }
}

fn main() {
    let chars = "hello".chars();
    chars.permutations(3).for_each(|cs| {
        println!("{:?}", cs);
    })
}