1use std::collections::BTreeMap;
2use std::collections::btree_map::IntoIter as BTreeIter;
3
4pub trait GroupByExt<I> {
5 fn group_by<K, F>(self, f: F) -> BTreeIter<K, Vec<I>> where
6 F: for<'a> Fn(&'a I) -> K,
7 K: Ord;
8}
9
10impl<T, I> GroupByExt<I> for T where T: Iterator<Item=I> {
11 fn group_by<K, F>(self, f: F) -> BTreeIter<K, Vec<I>> where
12 F: for<'a> Fn(&'a I) -> K,
13 K: Ord
14 {
15 let mut map = BTreeMap::new();
16 for item in self {
17 let hash = f(&item);
18 map.entry(hash).or_insert(vec![]).push(item);
19 }
20 map.into_iter()
21 }
22}
23
24#[test]
25fn test_group_by() {
26 let map: BTreeMap<u32, Vec<u32>> = (1..7).group_by(|i| i%3).collect();
27
28 assert_eq!(map[&0], vec![3, 6]);
29 assert_eq!(map[&1], vec![1, 4]);
30 assert_eq!(map[&2], vec![2, 5]);
31}