competitive_library/math/
permutation.rs

1//! 順列生成
2
3pub fn make_permutation(n: usize) -> Vec<Vec<usize>> {
4    let factorial = (1..=n).product();
5    let mut vvec: Vec<Vec<usize>> = vec![Vec::new(); factorial];
6    let nums: Vec<usize> = (0..n).collect();
7    let indexes: Vec<usize> = (0..factorial).collect();
8    push_recusive(nums, indexes, &mut vvec);
9    vvec
10}
11
12fn push_recusive<T: Clone>(
13    nums: Vec<T>,
14    indexes: Vec<usize>,
15    vvec: &mut Vec<Vec<T>>,
16) -> &mut Vec<Vec<T>> {
17    if nums.is_empty() {
18        return vvec;
19    }
20    let block_size = (1..nums.len()).product();
21    for (block_index, num) in nums.iter().enumerate() {
22        for inner_index in 0..block_size {
23            let index = indexes[block_size * block_index + inner_index];
24            vvec[index].push(num.clone());
25        }
26        let new_nums = {
27            let mut tmp = nums.clone();
28            tmp.remove(block_index);
29            tmp
30        };
31        let new_indexes: Vec<usize> = {
32            let slice = &indexes[(block_size * block_index)..(block_size * (block_index + 1))];
33            slice.to_vec()
34        };
35        push_recusive(new_nums, new_indexes, vvec);
36    }
37    vvec
38}
39
40pub struct Permutation<T>
41where
42    T: Clone,
43{
44    p: Vec<T>,
45    init: bool,
46}
47
48impl<T> Permutation<T>
49where
50    T: Clone,
51{
52    pub fn new(p: &[T]) -> Self {
53        Self {
54            p: p.to_vec(),
55            init: false,
56        }
57    }
58}
59impl<T> Iterator for Permutation<T>
60where
61    T: Clone + Ord,
62{
63    type Item = Vec<T>;
64
65    fn next(&mut self) -> Option<Self::Item> {
66        if !self.init {
67            self.p.sort();
68            self.init = true;
69            return Some(self.p.clone());
70        }
71        if self.p.len() < 2 {
72            return None;
73        }
74        let i = (0..self.p.len() - 1).rfind(|&i| self.p[i] < self.p[i + 1])?;
75        let j = self.p.iter().rposition(|x| x > &self.p[i]).unwrap();
76        self.p.swap(i, j);
77        self.p[i + 1..].reverse();
78        Some(self.p.clone())
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    #[test]
86    fn test_prm() {
87        let vv = make_permutation(4);
88        assert_eq!(0, vv[0][0]);
89    }
90    #[test]
91    fn test_struct() {
92        let expect = [
93            &[0, 1, 2],
94            &[0, 2, 1],
95            &[1, 0, 2],
96            &[1, 2, 0],
97            &[2, 0, 1],
98            &[2, 1, 0],
99        ];
100        let a = Permutation::new(&[0, 1, 2]);
101
102        for (i, v) in a.enumerate() {
103            assert_eq!(v, expect[i]);
104        }
105    }
106}