haar_lib/algo/
permutation.rs

1//! 順列の列挙
2use std::iter::{from_fn, once};
3
4macro_rules! impl_permutation {
5    ( $cmp1:tt, $cmp2:tt, $a: expr ) => {{
6        let n = $a.len();
7
8        if n <= 1 {
9            false
10        } else {
11            let i = (0..n - 1).rev().find(|&i| $a[i] $cmp1 $a[i + 1]);
12
13            match i {
14                None => false,
15                Some(i) => {
16                    let j = (i + 1..n).rev().find(|&j| $a[j] $cmp2 $a[i]).unwrap();
17
18                    $a.swap(i, j);
19                    $a[i + 1..].reverse();
20
21                    true
22                }
23            }
24        }
25    }}
26}
27
28/// `a`を辞書式順序で次の順列にする。
29pub fn next_permutation<T: Ord + Copy>(a: &mut [T]) -> bool {
30    impl_permutation!(<, >, a)
31}
32
33/// `a`を辞書式順序で前の順列にする。
34pub fn prev_permutation<T: Ord + Copy>(a: &mut [T]) -> bool {
35    impl_permutation!(>, <, a)
36}
37
38/// 辞書式順序で`a`以降の順列を列挙するイテレータを返す。
39pub fn permutations<T: Ord + Copy>(mut a: Vec<T>) -> impl Iterator<Item = Vec<T>> {
40    once(a.clone()).chain(from_fn(move || next_permutation(&mut a).then(|| a.clone())))
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    #[test]
48    fn test() {
49        let a = [1, 2, 3, 4, 5];
50
51        for a in permutations(a.to_vec()) {
52            let mut b = a.clone();
53            if next_permutation(&mut b) {
54                prev_permutation(&mut b);
55                assert_eq!(a, b);
56            }
57        }
58
59        let a = [1, 2, 3];
60
61        assert_eq!(
62            permutations(a.to_vec()).collect::<Vec<_>>(),
63            [
64                [1, 2, 3],
65                [1, 3, 2],
66                [2, 1, 3],
67                [2, 3, 1],
68                [3, 1, 2],
69                [3, 2, 1]
70            ]
71        );
72    }
73}