competitive_hpp/utils/
power_set.rs

1use std::collections::HashSet;
2
3pub trait PowerSetTrait<T> {
4    fn power_set(&self) -> Vec<Vec<T>>;
5}
6
7impl<T> PowerSetTrait<T> for Vec<T>
8where
9    T: Clone,
10{
11    #[inline]
12    fn power_set(&self) -> Vec<Vec<T>> {
13        let n = self.len();
14        (0..1 << n)
15            .map(|pattern| {
16                (0..n)
17                    .filter_map(|i| {
18                        if (pattern >> i) & 1 == 1 {
19                            Some(self[i].clone())
20                        } else {
21                            None
22                        }
23                    })
24                    .collect::<Vec<_>>()
25            })
26            .collect()
27    }
28}
29
30impl<K> PowerSetTrait<K> for HashSet<K>
31where
32    K: Clone,
33{
34    /// Unspecified order
35    #[inline]
36    fn power_set(&self) -> Vec<Vec<K>> {
37        let n = self.len();
38        let v: Vec<_> = self.iter().collect();
39        (0..1 << n)
40            .map(|pattern| {
41                (0..n)
42                    .filter_map(|i| {
43                        if (pattern >> i) & 1 == 1 {
44                            Some(v[i].clone())
45                        } else {
46                            None
47                        }
48                    })
49                    .collect::<Vec<_>>()
50            })
51            .collect()
52    }
53}
54
55#[cfg(test)]
56mod test {
57    use super::*;
58
59    #[test]
60    fn vec_test() {
61        let mut test = vec![0, 1, 2];
62
63        assert_eq!(
64            vec![
65                vec![],
66                vec![0],
67                vec![1],
68                vec![0, 1],
69                vec![2],
70                vec![0, 2],
71                vec![1, 2],
72                vec![0, 1, 2]
73            ],
74            test.power_set()
75        );
76
77        test.push(3);
78        test.push(4);
79
80        assert_eq!(test.power_set().len(), 32);
81    }
82
83    #[test]
84    fn set_test() {
85        let test = maplit::hashset! { 0, 1, 2 };
86
87        assert_eq!(
88            vec![
89                vec![],
90                vec![0],
91                vec![1],
92                vec![0, 1],
93                vec![2],
94                vec![0, 2],
95                vec![1, 2],
96                vec![0, 1, 2]
97            ]
98            .len(),
99            test.power_set().len()
100        );
101
102        let test_set = maplit::hashset! {
103            vec![],
104            vec![0],
105            vec![1],
106            vec![0, 1],
107            vec![2],
108            vec![0, 2],
109            vec![1, 2],
110            vec![0, 1, 2]
111        };
112        let test_power_set: HashSet<Vec<_>> = test.power_set().into_iter().collect();
113
114        assert_eq!(test_set.len(), test_power_set.len());
115        assert!(test_power_set.contains(&vec![0,]));
116    }
117}