Skip to main content

atcoder/lib/algorithm/
rle.rs

1//! ランレングス圧縮です。
2
3/// 連続する同一要素をまとめるイテレータ拡張です。
4///
5/// 要素同士を [`PartialEq`] で比較できるイテレータに利用できます。値の並べ替えは行わないため、
6/// 離れた位置にある同じ値は別々のグループになります。
7pub trait RleExt: Iterator {
8    /// イテレータを消費し、`(要素, 連続回数)` の列を返します。
9    ///
10    /// イテレータを先頭から末尾まで消費します。返り値は出現順で、各回数は 1 以上です。
11    /// 空イテレータでは空のベクタを返します。要素数を `N`、連続グループ数を `G` とすると、
12    /// 時間計算量は `O(N)`、結果に必要な追加メモリは `O(G)` です。
13    ///
14    /// # Examples
15    ///
16    /// ```
17    /// use atcoder::algorithm::rle::RleExt;
18    ///
19    /// assert_eq!("aaabbc".chars().rle(), vec![('a', 3), ('b', 2), ('c', 1)]);
20    /// ```
21    fn rle(self) -> Vec<(Self::Item, usize)>
22    where
23        Self: Sized,
24        Self::Item: PartialEq;
25}
26
27impl<I: Iterator> RleExt for I {
28    fn rle(self) -> Vec<(Self::Item, usize)>
29    where
30        Self: Sized,
31        Self::Item: PartialEq,
32    {
33        let mut iter = self;
34        let Some(mut current) = iter.next() else {
35            return Vec::new();
36        };
37        let mut count = 1;
38        let mut result = Vec::new();
39        for item in iter {
40            if item == current {
41                count += 1;
42            } else {
43                result.push((current, count));
44                current = item;
45                count = 1;
46            }
47        }
48        result.push((current, count));
49        result
50    }
51}