competitive/tools/
iterator_ext.rs

1use std::iter::Peekable;
2
3pub trait IteratorExt: Iterator {
4    fn merge_by<I, F>(self, other: I, is_first: F) -> MergeBy<Self, I, F>
5    where
6        Self: Sized,
7        I: Iterator<Item = Self::Item>,
8        F: FnMut(&Self::Item, &Self::Item) -> bool,
9    {
10        MergeBy {
11            left: self.peekable(),
12            right: other.peekable(),
13            is_first,
14        }
15    }
16}
17
18impl<I> IteratorExt for I where I: Iterator {}
19
20pub struct MergeBy<I, J, F>
21where
22    I: Iterator,
23    J: Iterator<Item = I::Item>,
24{
25    left: Peekable<I>,
26    right: Peekable<J>,
27    is_first: F,
28}
29
30impl<I, J, F> Iterator for MergeBy<I, J, F>
31where
32    I: Iterator,
33    J: Iterator<Item = I::Item>,
34    F: FnMut(&I::Item, &I::Item) -> bool,
35{
36    type Item = I::Item;
37
38    fn next(&mut self) -> Option<Self::Item> {
39        match (self.left.peek(), self.right.peek()) {
40            (Some(l), Some(r)) => {
41                if (self.is_first)(l, r) {
42                    self.left.next()
43                } else {
44                    self.right.next()
45                }
46            }
47            (Some(_), None) => self.left.next(),
48            (None, Some(_)) => self.right.next(),
49            (None, None) => None,
50        }
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    #[test]
59    fn test_merge_by() {
60        let a = vec![1, 4, 5];
61        let b = vec![2, 3, 6];
62        let merged: Vec<_> = a
63            .into_iter()
64            .merge_by(b.into_iter(), |x, y| x < y)
65            .collect();
66        assert_eq!(merged, vec![1, 2, 3, 4, 5, 6]);
67    }
68}
competitive/tools/iterator_ext.rs

competitive/tools/
iterator_ext.rs
