competitive/algorithm/
combinations.rs

1use std::{collections::BTreeSet, mem::swap};
2
3pub trait SliceCombinationsExt<T> {
4    fn for_each_product<F>(&self, r: usize, f: F)
5    where
6        F: FnMut(&[T]);
7    fn for_each_permutations<F>(&self, r: usize, f: F)
8    where
9        F: FnMut(&[T]);
10    fn for_each_combinations<F>(&self, r: usize, f: F)
11    where
12        F: FnMut(&[T]);
13    fn for_each_combinations_with_replacement<F>(&self, r: usize, f: F)
14    where
15        F: FnMut(&[T]);
16    fn next_permutation(&mut self) -> bool
17    where
18        T: Ord;
19    fn prev_permutation(&mut self) -> bool
20    where
21        T: Ord;
22    fn next_combination(&mut self, r: usize) -> bool
23    where
24        T: Ord;
25    fn prev_combination(&mut self, r: usize) -> bool
26    where
27        T: Ord;
28}
29
30impl<T> SliceCombinationsExt<T> for [T]
31where
32    T: Clone,
33{
34    /// choose `r` elements from `n` independently
35    ///
36    /// # Example
37    ///
38    /// ```
39    /// # use competitive::algorithm::SliceCombinationsExt;
40    /// let n = vec![1, 2, 3, 4];
41    /// let mut p = Vec::new();
42    /// let mut q = Vec::new();
43    /// n.for_each_product(2, |v| p.push(v.to_vec()));
44    /// for x in n.iter().cloned() {
45    ///     for y in n.iter().cloned() {
46    ///         q.push(vec![x, y]);
47    ///     }
48    /// }
49    /// assert_eq!(p, q);
50    /// ```
51    fn for_each_product<F>(&self, r: usize, mut f: F)
52    where
53        F: FnMut(&[T]),
54    {
55        fn product_inner<T, F>(n: &[T], mut r: usize, buf: &mut Vec<T>, f: &mut F)
56        where
57            T: Clone,
58            F: FnMut(&[T]),
59        {
60            if r == 0 {
61                f(buf.as_slice());
62            } else {
63                r -= 1;
64                for a in n.iter().cloned() {
65                    buf.push(a);
66                    product_inner(n, r, buf, f);
67                    buf.pop();
68                }
69            }
70        }
71
72        let mut v = Vec::with_capacity(r);
73        product_inner(self, r, &mut v, &mut f);
74    }
75
76    /// choose `r` elements from `n` independently
77    ///
78    /// # Example
79    ///
80    /// ```
81    /// # use competitive::algorithm::SliceCombinationsExt;
82    /// let n = vec![1, 2, 3, 4];
83    /// let mut p = Vec::new();
84    /// let mut q = Vec::new();
85    /// n.for_each_product(2, |v| p.push(v.to_vec()));
86    /// for x in n.iter().cloned() {
87    ///     for y in n.iter().cloned() {
88    ///         q.push(vec![x, y]);
89    ///     }
90    /// }
91    /// assert_eq!(p, q);
92    /// ```
93    fn for_each_permutations<F>(&self, r: usize, mut f: F)
94    where
95        F: FnMut(&[T]),
96    {
97        fn permutations_inner<T, F>(
98            n: &[T],
99            mut r: usize,
100            rem: &mut BTreeSet<usize>,
101            buf: &mut Vec<T>,
102            f: &mut F,
103        ) where
104            T: Clone,
105            F: FnMut(&[T]),
106        {
107            if r == 0 {
108                f(buf.as_slice());
109            } else {
110                r -= 1;
111                for i in rem.iter().cloned().collect::<Vec<_>>() {
112                    buf.push(n[i].clone());
113                    rem.remove(&i);
114                    permutations_inner(n, r, rem, buf, f);
115                    rem.insert(i);
116                    buf.pop();
117                }
118            }
119        }
120
121        if r <= self.len() {
122            let mut v = Vec::with_capacity(r);
123            let mut rem: BTreeSet<usize> = (0..self.len()).collect();
124            permutations_inner(self, r, &mut rem, &mut v, &mut f);
125        }
126    }
127
128    /// choose distinct `r` elements from `n` in any order
129    ///
130    /// # Example
131    ///
132    /// ```
133    /// # use competitive::algorithm::SliceCombinationsExt;
134    /// let n = vec![1, 2, 3, 4];
135    /// let mut p = Vec::new();
136    /// let mut q = Vec::new();
137    /// n.for_each_permutations(2, |v| p.push(v.to_vec()));
138    /// for (i, x) in n.iter().cloned().enumerate() {
139    ///     for (j, y) in n.iter().cloned().enumerate() {
140    ///         if i != j {
141    ///             q.push(vec![x, y]);
142    ///         }
143    ///     }
144    /// }
145    /// assert_eq!(p, q);
146    /// ```
147    fn for_each_combinations<F>(&self, r: usize, mut f: F)
148    where
149        F: FnMut(&[T]),
150    {
151        fn combinations_inner<T, F>(
152            n: &[T],
153            mut r: usize,
154            start: usize,
155            buf: &mut Vec<T>,
156            f: &mut F,
157        ) where
158            T: Clone,
159            F: FnMut(&[T]),
160        {
161            if r == 0 {
162                f(buf.as_slice());
163            } else {
164                r -= 1;
165                for i in start..n.len() - r {
166                    buf.push(n[i].clone());
167                    combinations_inner(n, r, i + 1, buf, f);
168                    buf.pop();
169                }
170            }
171        }
172
173        if r <= self.len() {
174            let mut v = Vec::with_capacity(r);
175            combinations_inner(self, r, 0, &mut v, &mut f);
176        }
177    }
178
179    /// choose `r` elements from `n` in sorted order
180    ///
181    /// # Example
182    ///
183    /// ```
184    /// # use competitive::algorithm::SliceCombinationsExt;
185    /// let n = vec![1, 2, 3, 4];
186    /// let mut p = Vec::new();
187    /// let mut q = Vec::new();
188    /// n.for_each_combinations_with_replacement(2, |v| p.push(v.to_vec()));
189    /// for (i, x) in n.iter().cloned().enumerate() {
190    ///     for y in n[i..].iter().cloned() {
191    ///         q.push(vec![x, y]);
192    ///     }
193    /// }
194    /// assert_eq!(p, q);
195    /// ```
196    fn for_each_combinations_with_replacement<F>(&self, r: usize, mut f: F)
197    where
198        F: FnMut(&[T]),
199    {
200        fn combinations_with_replacement_inner<T, F>(
201            n: &[T],
202            mut r: usize,
203            start: usize,
204            buf: &mut Vec<T>,
205            f: &mut F,
206        ) where
207            T: Clone,
208            F: FnMut(&[T]),
209        {
210            if r == 0 {
211                f(buf.as_slice());
212            } else {
213                r -= 1;
214                for i in start..n.len() {
215                    buf.push(n[i].clone());
216                    combinations_with_replacement_inner(n, r, i, buf, f);
217                    buf.pop();
218                }
219            }
220        }
221
222        if r <= self.len() {
223            let mut v = Vec::with_capacity(r);
224            combinations_with_replacement_inner(self, r, 0, &mut v, &mut f);
225        }
226    }
227
228    /// Permute the elements into next permutation in lexicographical order.
229    /// Return whether such a next permutation exists.
230    fn next_permutation(&mut self) -> bool
231    where
232        T: Ord,
233    {
234        if self.len() < 2 {
235            return false;
236        }
237        let mut target = self.len() - 2;
238        while target > 0 && self[target] > self[target + 1] {
239            target -= 1;
240        }
241        if target == 0 && self[target] > self[target + 1] {
242            return false;
243        }
244        let mut next = self.len() - 1;
245        while next > target && self[next] < self[target] {
246            next -= 1;
247        }
248        self.swap(next, target);
249        self[target + 1..].reverse();
250        true
251    }
252
253    /// Permute the elements into previous permutation in lexicographical order.
254    /// Return whether such a previous permutation exists.
255    fn prev_permutation(&mut self) -> bool
256    where
257        T: Ord,
258    {
259        if self.len() < 2 {
260            return false;
261        }
262        let mut target = self.len() - 2;
263        while target > 0 && self[target] < self[target + 1] {
264            target -= 1;
265        }
266        if target == 0 && self[target] < self[target + 1] {
267            return false;
268        }
269        self[target + 1..].reverse();
270        let mut next = self.len() - 1;
271        while next > target && self[next - 1] < self[target] {
272            next -= 1;
273        }
274        self.swap(target, next);
275        true
276    }
277
278    /// Permute the elements into next combination choosing r elements in lexicographical order.
279    /// Return whether such a next combination exists.
280    fn next_combination(&mut self, r: usize) -> bool
281    where
282        T: Ord,
283    {
284        assert!(r <= self.len());
285        let (a, b) = self.split_at_mut(r);
286        next_combination_inner(a, b)
287    }
288
289    /// Permute the elements into previous combination choosing r elements in lexicographical order.
290    /// Return whether such a previous combination exists.
291    fn prev_combination(&mut self, r: usize) -> bool
292    where
293        T: Ord,
294    {
295        assert!(r <= self.len());
296        let (a, b) = self.split_at_mut(r);
297        next_combination_inner(b, a)
298    }
299}
300
301fn rotate_distinct<'a, T>(mut a: &'a mut [T], mut b: &'a mut [T]) {
302    while !a.is_empty() && !b.is_empty() {
303        if a.len() >= b.len() {
304            let (l, r) = a.split_at_mut(b.len());
305            l.swap_with_slice(b);
306            a = r;
307        } else {
308            let (l, r) = b.split_at_mut(a.len());
309            l.swap_with_slice(a);
310            a = l;
311            b = r;
312        }
313    }
314}
315
316fn next_combination_inner<T>(a: &mut [T], b: &mut [T]) -> bool
317where
318    T: Ord,
319{
320    if a.is_empty() || b.is_empty() {
321        return false;
322    }
323    let mut target = a.len() - 1;
324    let last_elem = b.last().unwrap();
325    while target > 0 && &a[target] >= last_elem {
326        target -= 1;
327    }
328    if target == 0 && &a[target] >= last_elem {
329        rotate_distinct(a, b);
330        return false;
331    }
332    let mut next = 0;
333    while a[target] >= b[next] {
334        next += 1;
335    }
336    swap(&mut a[target], &mut b[next]);
337    rotate_distinct(&mut a[target + 1..], &mut b[next + 1..]);
338    true
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344
345    #[test]
346    fn test_next_prev_permutation() {
347        for n in 1..=7 {
348            let mut p: Vec<_> = (0..n).collect();
349            let mut a = vec![];
350            p.for_each_permutations(n, |p| a.push(p.to_vec()));
351            let mut b = vec![];
352            loop {
353                b.push(p.to_vec());
354                if !p.next_permutation() {
355                    break;
356                }
357                assert!(p.prev_permutation());
358                assert_eq!(b.last().as_ref().unwrap().as_slice(), &p);
359                assert!(p.next_permutation());
360            }
361            assert_eq!(a, b);
362        }
363    }
364
365    #[test]
366    fn test_next_prev_combination() {
367        for n in 1..=7 {
368            for r in 0..=n {
369                let mut p: Vec<_> = (0..n).collect();
370                let mut a = vec![];
371                p.for_each_combinations(r, |p| a.push(p.to_vec()));
372                let mut b = vec![];
373                loop {
374                    b.push(p[..r].to_vec());
375                    if !p.next_combination(r) {
376                        break;
377                    }
378                    assert!(p.prev_combination(r));
379                    assert_eq!(b.last().as_ref().unwrap().as_slice(), &p[..r]);
380                    assert!(p.next_combination(r));
381                }
382                assert_eq!(a, b);
383            }
384        }
385    }
386}
competitive/algorithm/combinations.rs

competitive/algorithm/
combinations.rs
