partition_iterator/
partit.rs

1use std::cmp::min;
2
3use crate::seqit::IncSeqIterator;
4
5pub struct Partitions<I>
6    where I: Iterator {
7    n: u32,
8    k: u32,
9    data: Vec<I::Item>,
10    seqs: IncSeqIterator
11}
12
13pub struct KPartitions<I> 
14    where I: Iterator   {
15    n: u32,
16    k: u32,
17    data: Vec<I::Item>,
18    seqs: IncSeqIterator
19}
20
21impl<I,T> Iterator for KPartitions<I>
22    where I: Iterator<Item = T>,
23          T: Clone
24{
25    type Item = Vec<Vec<T>>;
26
27    fn next(&mut self) -> Option<Self::Item> {
28        match self.seqs.next() {
29            Some(seq) => {
30                let mut res:Self::Item = vec![vec![]; self.k as usize];
31                for i in 0..(self.n as usize) {
32                    res[seq[i] as usize].push(self.data[i].clone());
33                }
34                Some(res)
35            }
36            None => None,
37        }
38    }
39}
40
41
42impl<I,T> Iterator for Partitions<I>
43    where I: Iterator<Item = T>,
44          T: Clone
45{
46    type Item = Vec<Vec<T>>;
47
48    fn next(&mut self) -> Option<Self::Item> {
49        loop {
50            if let Some(seq) = self.seqs.next() {
51                let mut res:Self::Item = vec![vec![]; self.k as usize];
52                for i in 0..(self.n as usize) {
53                    res[seq[i] as usize].push(self.data[i].clone());
54                }
55                return Some(res)
56            }
57
58            self.k += 1;
59            if self.k > self.n {
60                return None
61            }
62            self.seqs = IncSeqIterator::new(self.n, self.k);
63        }
64    }
65}
66
67pub trait PartitionIter<I> where I: Iterator {
68    fn kpartitions(self, k:u32) -> KPartitions<I>;
69    fn partitions(self) -> Partitions<I>;
70}
71
72impl<I> PartitionIter<I> for I where I: Iterator + Sized
73{
74    fn kpartitions(self, k:u32) -> KPartitions<I> {
75        let data:Vec<I::Item> = self.collect(); // TODO: make lazy
76        let n = data.len() as u32;
77        KPartitions{ n, k, data, seqs: IncSeqIterator::new(n, k) }
78    }
79
80    fn partitions(self) -> Partitions<I> {
81        let data:Vec<I::Item> = self.collect(); // TODO: make lazy
82        let n = data.len() as u32;
83        let k = min(1, n); // k = 0 makes only sense for n = 0
84        Partitions{ n, k, data, seqs: IncSeqIterator::new(n, k) }
85    }    
86}
87
88
89