cp_library_rs/utils/
lineartime_merging.rs

1//! ソート済み配列の線形時間マージ
2
3use std::iter::Peekable;
4
5pub trait Merge: Iterator
6where
7    Self::Item: Ord,
8    Self: Sized,
9{
10    /// ソート済み配列をマージする
11    ///
12    /// - 時間計算量: $`O(N + M)`$
13    fn merge_linear(self, other: Self) -> MergeIterator<Self::Item, Self, Self> {
14        MergeIterator {
15            itr_a: self.peekable(),
16            itr_b: other.peekable(),
17        }
18    }
19}
20
21impl<I: Iterator> Merge for I where I::Item: Ord {}
22
23/// 2つのソート済みイテレータを,整列順に返す
24pub struct MergeIterator<T, I, J>
25where
26    T: Ord,
27    I: Iterator<Item = T>,
28    J: Iterator<Item = T>,
29{
30    itr_a: Peekable<I>,
31    itr_b: Peekable<J>,
32}
33
34impl<T, I, J> Iterator for MergeIterator<T, I, J>
35where
36    T: Ord,
37    I: Iterator<Item = T>,
38    J: Iterator<Item = T>,
39{
40    type Item = T;
41    fn next(&mut self) -> Option<Self::Item> {
42        match (self.itr_a.peek(), self.itr_b.peek()) {
43            (Some(a), Some(b)) => {
44                if a <= b {
45                    self.itr_a.next()
46                } else {
47                    self.itr_b.next()
48                }
49            }
50            (Some(_), _) => self.itr_a.next(),
51            _ => self.itr_b.next(),
52        }
53    }
54}