Skip to main content

routee_compass_core/util/
multiset.rs

1pub struct MultiSet<'a, T: Clone + Copy> {
2    sets: &'a Vec<Vec<T>>,
3    pos: Option<Vec<usize>>,
4    final_pos: Vec<usize>,
5}
6
7impl<'a, T> From<&'a Vec<Vec<T>>> for MultiSet<'a, T>
8where
9    T: Clone + Copy,
10{
11    fn from(sets: &'a Vec<Vec<T>>) -> Self {
12        let final_pos: Vec<usize> = sets.iter().map(|v| v.len() - 1).collect();
13        let pos: Option<Vec<usize>> = Some(vec![0; sets.len()]);
14        MultiSet {
15            sets,
16            pos,
17            final_pos,
18        }
19    }
20}
21
22impl<T> Iterator for MultiSet<'_, T>
23where
24    T: Clone + Copy,
25{
26    type Item = Vec<T>;
27
28    fn next(&mut self) -> Option<Self::Item> {
29        match &self.pos {
30            None => None,
31            Some(position) => {
32                let result: Vec<T> = position
33                    .iter()
34                    .zip(0..self.sets.len())
35                    .map(|(j, i)| self.sets[i][*j])
36                    .collect();
37                let mut next_pos = position.to_vec();
38                // tick through full table of index combinations
39                let mut finished = false;
40                for idx in 0..self.sets.len() {
41                    if next_pos[idx] < self.final_pos[idx] {
42                        next_pos[idx] += 1;
43                        break;
44                    } else if idx == self.sets.len() - 1 {
45                        finished = true;
46                        break;
47                    } else {
48                        for r in next_pos.iter_mut().take(idx + 1) {
49                            *r = 0;
50                        }
51                    }
52                }
53                if finished {
54                    self.pos = None;
55                } else {
56                    self.pos = Some(next_pos);
57                }
58                Some(result)
59            }
60        }
61    }
62}
63
64#[cfg(test)]
65mod test {
66
67    use super::MultiSet;
68
69    #[test]
70    fn test_multiset() {
71        let input = vec![vec![1, 3], vec![2], vec![5, 7, 9]];
72        let ms = MultiSet::from(&input);
73        let result: Vec<Vec<i32>> = ms.into_iter().collect();
74        let expected = [
75            [1, 2, 5],
76            [3, 2, 5],
77            [1, 2, 7],
78            [3, 2, 7],
79            [1, 2, 9],
80            [3, 2, 9],
81        ];
82        assert_eq!(result, expected);
83    }
84}