nekolib/algo/
rle.rs

1pub mod rle {
2    pub struct Rle<T, I> {
3        iter: I,
4        last: Option<T>,
5    }
6    pub struct RleBy<T, I, F> {
7        iter: I,
8        last: Option<T>,
9        compare: F,
10    }
11    pub struct RleByKey<T, I, K> {
12        iter: I,
13        last: Option<T>,
14        key: K,
15    }
16
17    impl<T: PartialEq, I: Iterator<Item = T>> Rle<T, I> {
18        pub fn new(iter: I) -> Self { Self { iter, last: None } }
19    }
20    impl<T, I: Iterator<Item = T>, F: FnMut(&T, &T) -> bool> RleBy<T, I, F> {
21        pub fn new(iter: I, compare: F) -> Self {
22            Self { iter, last: None, compare }
23        }
24    }
25    impl<T, I: Iterator<Item = T>, U: PartialEq, K: FnMut(&T) -> U>
26        RleByKey<T, I, K>
27    {
28        pub fn new(iter: I, key: K) -> Self { Self { iter, last: None, key } }
29    }
30
31    impl<T: PartialEq, I: Iterator<Item = T>> Iterator for Rle<T, I> {
32        type Item = (usize, T);
33
34        fn next(&mut self) -> Option<(usize, T)> {
35            let last = self.last.take().or_else(|| self.iter.next())?;
36            let mut count = 1;
37            while let Some(next) = self.iter.next() {
38                if last == next {
39                    count += 1;
40                } else {
41                    self.last = Some(next);
42                    return Some((count, last));
43                }
44            }
45            Some((count, last))
46        }
47    }
48
49    impl<T, I: Iterator<Item = T>, F: FnMut(&T, &T) -> bool> Iterator
50        for RleBy<T, I, F>
51    {
52        type Item = (usize, T);
53
54        fn next(&mut self) -> Option<(usize, T)> {
55            let last = self.last.take().or_else(|| self.iter.next())?;
56            let mut count = 1;
57            while let Some(next) = self.iter.next() {
58                if (self.compare)(&last, &next) {
59                    count += 1;
60                } else {
61                    self.last = Some(next);
62                    return Some((count, last));
63                }
64            }
65            Some((count, last))
66        }
67    }
68
69    impl<T, I: Iterator<Item = T>, U: PartialEq, K: FnMut(&T) -> U> Iterator
70        for RleByKey<T, I, K>
71    {
72        type Item = (usize, T);
73
74        fn next(&mut self) -> Option<(usize, T)> {
75            let last = self.last.take().or_else(|| self.iter.next())?;
76            let mut count = 1;
77            while let Some(next) = self.iter.next() {
78                if (self.key)(&last) == (self.key)(&next) {
79                    count += 1;
80                } else {
81                    self.last = Some(next);
82                    return Some((count, last));
83                }
84            }
85            Some((count, last))
86        }
87    }
88}
89
90pub trait Rle<T, I> {
91    fn rle(self) -> rle::Rle<T, I>;
92}
93
94pub trait RleBy<T, I> {
95    fn rle_by<F: FnMut(&T, &T) -> bool>(
96        self,
97        compare: F,
98    ) -> rle::RleBy<T, I, F>;
99}
100
101pub trait RleByKey<T, I> {
102    fn rle_by_key<U: PartialEq, K: FnMut(&T) -> U>(
103        self,
104        key: K,
105    ) -> rle::RleByKey<T, I, K>;
106}
107
108impl<T: PartialEq, I: Iterator<Item = T>> Rle<T, I> for I {
109    fn rle(self) -> rle::Rle<T, I> { rle::Rle::new(self) }
110}
111
112impl<T, I: Iterator<Item = T>> RleBy<T, I> for I {
113    fn rle_by<F: FnMut(&T, &T) -> bool>(
114        self,
115        compare: F,
116    ) -> rle::RleBy<T, I, F> {
117        rle::RleBy::new(self, compare)
118    }
119}
120
121impl<T, I: Iterator<Item = T>> RleByKey<T, I> for I {
122    fn rle_by_key<U: PartialEq, K: FnMut(&T) -> U>(
123        self,
124        key: K,
125    ) -> rle::RleByKey<T, I, K> {
126        rle::RleByKey::new(self, key)
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use crate::{Rle, RleByKey};
133
134    #[test]
135    fn sanity_check() {
136        let a = vec![1, 1, 3, 1, 2, 5, 5, 5, 5, 3];
137        let rle: Vec<_> = a.into_iter().rle().collect();
138        assert_eq!(rle, [(2, 1), (1, 3), (1, 1), (1, 2), (4, 5), (1, 3)]);
139    }
140
141    #[test]
142    fn empty() {
143        assert!(std::iter::empty::<()>().rle().next().is_none());
144    }
145
146    #[test]
147    fn by_key() {
148        let a = vec![1, 1, 3, 1, 2, 5, 5, 5, 5, 3];
149        let rle: Vec<_> =
150            a.into_iter().enumerate().rle_by_key(|&(_i, x)| x).collect();
151        assert_eq!(rle, [
152            (2, (0, 1)),
153            (1, (2, 3)),
154            (1, (3, 1)),
155            (1, (4, 2)),
156            (4, (5, 5)),
157            (1, (9, 3)),
158        ]);
159    }
160}
nekolib/algo/rle.rs

nekolib/algo/
rle.rs
