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}