Skip to main content

nekolib/traits/
group_by.rs

1//! イテレータのグルーピング。
2
3use std::collections::BTreeMap;
4
5/// イテレータのグルーピング。
6///
7/// # Suggestions
8/// `BTreeMap` と `HashMap` 用で分けるべき? メソッド名を冗長にしたくない。
9///
10/// # See also
11/// [`UsizeGroupBy`]
12///
13/// [`UsizeGroupBy`]: trait.UsizeGroupBy.html
14pub trait GroupBy<V> {
15    /// # Examples
16    /// ```
17    /// use nekolib::traits::GroupBy;
18    ///
19    /// let a = vec![1, 4, 3, -5, -6, 0, 2, -2, 3];
20    /// let g1 = a.iter().copied().group_by(|&ai: &i32| ai.rem_euclid(3));
21    /// let g2 = a.iter().copied().group_by(|&ai: &i32| ai % 3);
22    ///
23    /// assert_eq!(g1.len(), 3);
24    /// assert_eq!(g1[&0], [3, -6, 0, 3]);
25    /// assert_eq!(g1[&1], [1, 4, -5, -2]);
26    /// assert_eq!(g1[&2], [2]);
27    ///
28    /// assert_eq!(g2.len(), 4);
29    /// assert_eq!(g2[&-2], [-5, -2]);
30    /// assert_eq!(g2[&0], [3, -6, 0, 3]);
31    /// assert_eq!(g2[&1], [1, 4]);
32    /// assert_eq!(g2[&2], [2]);
33    /// ```
34    fn group_by<K: Ord>(self, key: impl FnMut(&V) -> K) -> BTreeMap<K, Vec<V>>;
35}
36
37impl<V, I: Iterator<Item = V>> GroupBy<V> for I {
38    fn group_by<K: Ord>(
39        self,
40        mut key: impl FnMut(&V) -> K,
41    ) -> BTreeMap<K, Vec<V>> {
42        let mut res = BTreeMap::new();
43        for v in self {
44            let k = key(&v);
45            res.entry(k).or_insert(vec![]).push(v);
46        }
47        res
48    }
49}