1use crate::{
2 adaptors::map::{MapSpecialCase, MapSpecialCaseFn},
3 MinMaxResult,
4};
5use std::cmp::Ordering;
6use std::collections::HashMap;
7use std::hash::Hash;
8use std::iter::Iterator;
9use std::ops::{Add, Mul};
10
11/// A wrapper to allow for an easy [`into_grouping_map_by`](crate::Itertools::into_grouping_map_by)
12pub type MapForGrouping<I, F> = MapSpecialCase<I, GroupingMapFn<F>>;
13
14#[derive(Clone)]
15pub struct GroupingMapFn<F>(F);
16
17impl<F> std::fmt::Debug for GroupingMapFn<F> {
18 debug_fmt_fields!(GroupingMapFn,);
19}
20
21impl<V, K, F: FnMut(&V) -> K> MapSpecialCaseFn<V> for GroupingMapFn<F> {
22 type Out = (K, V);
23 fn call(&mut self, v: V) -> Self::Out {
24 ((self.0)(&v), v)
25 }
26}
27
28pub(crate) fn new_map_for_grouping<K, I: Iterator, F: FnMut(&I::Item) -> K>(
29 iter: I,
30 key_mapper: F,
31) -> MapForGrouping<I, F> {
32 MapSpecialCase {
33 iter,
34 f: GroupingMapFn(key_mapper),
35 }
36}
37
38/// Creates a new `GroupingMap` from `iter`
39pub fn new<I, K, V>(iter: I) -> GroupingMap<I>
40where
41 I: Iterator<Item = (K, V)>,
42 K: Hash + Eq,
43{
44 GroupingMap { iter }
45}
46
47/// `GroupingMapBy` is an intermediate struct for efficient group-and-fold operations.
48///
49/// See [`GroupingMap`] for more informations.
50pub type GroupingMapBy<I, F> = GroupingMap<MapForGrouping<I, F>>;
51
52/// `GroupingMap` is an intermediate struct for efficient group-and-fold operations.
53/// It groups elements by their key and at the same time fold each group
54/// using some aggregating operation.
55///
56/// No method on this struct performs temporary allocations.
57#[derive(Clone, Debug)]
58#[must_use = "GroupingMap is lazy and do nothing unless consumed"]
59pub struct GroupingMap<I> {
60 iter: I,
61}
62
63impl<I, K, V> GroupingMap<I>
64where
65 I: Iterator<Item = (K, V)>,
66 K: Hash + Eq,
67{
68 /// This is the generic way to perform any operation on a `GroupingMap`.
69 /// It's suggested to use this method only to implement custom operations
70 /// when the already provided ones are not enough.
71 ///
72 /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
73 /// of each group sequentially, passing the previously accumulated value, a reference to the key
74 /// and the current element as arguments, and stores the results in an `HashMap`.
75 ///
76 /// The `operation` function is invoked on each element with the following parameters:
77 /// - the current value of the accumulator of the group if there is currently one;
78 /// - a reference to the key of the group this element belongs to;
79 /// - the element from the source being aggregated;
80 ///
81 /// If `operation` returns `Some(element)` then the accumulator is updated with `element`,
82 /// otherwise the previous accumulation is discarded.
83 ///
84 /// Return a `HashMap` associating the key of each group with the result of aggregation of
85 /// that group's elements. If the aggregation of the last element of a group discards the
86 /// accumulator then there won't be an entry associated to that group's key.
87 ///
88 /// ```
89 /// use itertools::Itertools;
90 ///
91 /// let data = vec![2, 8, 5, 7, 9, 0, 4, 10];
92 /// let lookup = data.into_iter()
93 /// .into_grouping_map_by(|&n| n % 4)
94 /// .aggregate(|acc, _key, val| {
95 /// if val == 0 || val == 10 {
96 /// None
97 /// } else {
98 /// Some(acc.unwrap_or(0) + val)
99 /// }
100 /// });
101 ///
102 /// assert_eq!(lookup[&0], 4); // 0 resets the accumulator so only 4 is summed
103 /// assert_eq!(lookup[&1], 5 + 9);
104 /// assert_eq!(lookup.get(&2), None); // 10 resets the accumulator and nothing is summed afterward
105 /// assert_eq!(lookup[&3], 7);
106 /// assert_eq!(lookup.len(), 3); // The final keys are only 0, 1 and 2
107 /// ```
108 pub fn aggregate<FO, R>(self, mut operation: FO) -> HashMap<K, R>
109 where
110 FO: FnMut(Option<R>, &K, V) -> Option<R>,
111 {
112 let mut destination_map = HashMap::new();
113
114 self.iter.for_each(|(key, val)| {
115 let acc = destination_map.remove(&key);
116 if let Some(op_res) = operation(acc, &key, val) {
117 destination_map.insert(key, op_res);
118 }
119 });
120
121 destination_map
122 }
123
124 /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
125 /// of each group sequentially, passing the previously accumulated value, a reference to the key
126 /// and the current element as arguments, and stores the results in a new map.
127 ///
128 /// `init` is called to obtain the initial value of each accumulator.
129 ///
130 /// `operation` is a function that is invoked on each element with the following parameters:
131 /// - the current value of the accumulator of the group;
132 /// - a reference to the key of the group this element belongs to;
133 /// - the element from the source being accumulated.
134 ///
135 /// Return a `HashMap` associating the key of each group with the result of folding that group's elements.
136 ///
137 /// ```
138 /// use itertools::Itertools;
139 ///
140 /// #[derive(Debug, Default)]
141 /// struct Accumulator {
142 /// acc: usize,
143 /// }
144 ///
145 /// let lookup = (1..=7)
146 /// .into_grouping_map_by(|&n| n % 3)
147 /// .fold_with(|_key, _val| Default::default(), |Accumulator { acc }, _key, val| {
148 /// let acc = acc + val;
149 /// Accumulator { acc }
150 /// });
151 ///
152 /// assert_eq!(lookup[&0].acc, 3 + 6);
153 /// assert_eq!(lookup[&1].acc, 1 + 4 + 7);
154 /// assert_eq!(lookup[&2].acc, 2 + 5);
155 /// assert_eq!(lookup.len(), 3);
156 /// ```
157 pub fn fold_with<FI, FO, R>(self, mut init: FI, mut operation: FO) -> HashMap<K, R>
158 where
159 FI: FnMut(&K, &V) -> R,
160 FO: FnMut(R, &K, V) -> R,
161 {
162 self.aggregate(|acc, key, val| {
163 let acc = acc.unwrap_or_else(|| init(key, &val));
164 Some(operation(acc, key, val))
165 })
166 }
167
168 /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
169 /// of each group sequentially, passing the previously accumulated value, a reference to the key
170 /// and the current element as arguments, and stores the results in a new map.
171 ///
172 /// `init` is the value from which will be cloned the initial value of each accumulator.
173 ///
174 /// `operation` is a function that is invoked on each element with the following parameters:
175 /// - the current value of the accumulator of the group;
176 /// - a reference to the key of the group this element belongs to;
177 /// - the element from the source being accumulated.
178 ///
179 /// Return a `HashMap` associating the key of each group with the result of folding that group's elements.
180 ///
181 /// ```
182 /// use itertools::Itertools;
183 ///
184 /// let lookup = (1..=7)
185 /// .into_grouping_map_by(|&n| n % 3)
186 /// .fold(0, |acc, _key, val| acc + val);
187 ///
188 /// assert_eq!(lookup[&0], 3 + 6);
189 /// assert_eq!(lookup[&1], 1 + 4 + 7);
190 /// assert_eq!(lookup[&2], 2 + 5);
191 /// assert_eq!(lookup.len(), 3);
192 /// ```
193 pub fn fold<FO, R>(self, init: R, operation: FO) -> HashMap<K, R>
194 where
195 R: Clone,
196 FO: FnMut(R, &K, V) -> R,
197 {
198 self.fold_with(|_, _| init.clone(), operation)
199 }
200
201 /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
202 /// of each group sequentially, passing the previously accumulated value, a reference to the key
203 /// and the current element as arguments, and stores the results in a new map.
204 ///
205 /// This is similar to [`fold`] but the initial value of the accumulator is the first element of the group.
206 ///
207 /// `operation` is a function that is invoked on each element with the following parameters:
208 /// - the current value of the accumulator of the group;
209 /// - a reference to the key of the group this element belongs to;
210 /// - the element from the source being accumulated.
211 ///
212 /// Return a `HashMap` associating the key of each group with the result of folding that group's elements.
213 ///
214 /// [`fold`]: GroupingMap::fold
215 ///
216 /// ```
217 /// use itertools::Itertools;
218 ///
219 /// let lookup = (1..=7)
220 /// .into_grouping_map_by(|&n| n % 3)
221 /// .reduce(|acc, _key, val| acc + val);
222 ///
223 /// assert_eq!(lookup[&0], 3 + 6);
224 /// assert_eq!(lookup[&1], 1 + 4 + 7);
225 /// assert_eq!(lookup[&2], 2 + 5);
226 /// assert_eq!(lookup.len(), 3);
227 /// ```
228 pub fn reduce<FO>(self, mut operation: FO) -> HashMap<K, V>
229 where
230 FO: FnMut(V, &K, V) -> V,
231 {
232 self.aggregate(|acc, key, val| {
233 Some(match acc {
234 Some(acc) => operation(acc, key, val),
235 None => val,
236 })
237 })
238 }
239
240 /// See [`.reduce()`](GroupingMap::reduce).
241 #[deprecated(note = "Use .reduce() instead", since = "0.13.0")]
242 pub fn fold_first<FO>(self, operation: FO) -> HashMap<K, V>
243 where
244 FO: FnMut(V, &K, V) -> V,
245 {
246 self.reduce(operation)
247 }
248
249 /// Groups elements from the `GroupingMap` source by key and collects the elements of each group in
250 /// an instance of `C`. The iteration order is preserved when inserting elements.
251 ///
252 /// Return a `HashMap` associating the key of each group with the collection containing that group's elements.
253 ///
254 /// ```
255 /// use itertools::Itertools;
256 /// use std::collections::HashSet;
257 ///
258 /// let lookup = vec![0, 1, 2, 3, 4, 5, 6, 2, 3, 6].into_iter()
259 /// .into_grouping_map_by(|&n| n % 3)
260 /// .collect::<HashSet<_>>();
261 ///
262 /// assert_eq!(lookup[&0], vec![0, 3, 6].into_iter().collect::<HashSet<_>>());
263 /// assert_eq!(lookup[&1], vec![1, 4].into_iter().collect::<HashSet<_>>());
264 /// assert_eq!(lookup[&2], vec![2, 5].into_iter().collect::<HashSet<_>>());
265 /// assert_eq!(lookup.len(), 3);
266 /// ```
267 pub fn collect<C>(self) -> HashMap<K, C>
268 where
269 C: Default + Extend<V>,
270 {
271 let mut destination_map = HashMap::new();
272
273 self.iter.for_each(|(key, val)| {
274 destination_map
275 .entry(key)
276 .or_insert_with(C::default)
277 .extend(Some(val));
278 });
279
280 destination_map
281 }
282
283 /// Groups elements from the `GroupingMap` source by key and finds the maximum of each group.
284 ///
285 /// If several elements are equally maximum, the last element is picked.
286 ///
287 /// Returns a `HashMap` associating the key of each group with the maximum of that group's elements.
288 ///
289 /// ```
290 /// use itertools::Itertools;
291 ///
292 /// let lookup = vec![1, 3, 4, 5, 7, 8, 9, 12].into_iter()
293 /// .into_grouping_map_by(|&n| n % 3)
294 /// .max();
295 ///
296 /// assert_eq!(lookup[&0], 12);
297 /// assert_eq!(lookup[&1], 7);
298 /// assert_eq!(lookup[&2], 8);
299 /// assert_eq!(lookup.len(), 3);
300 /// ```
301 pub fn max(self) -> HashMap<K, V>
302 where
303 V: Ord,
304 {
305 self.max_by(|_, v1, v2| V::cmp(v1, v2))
306 }
307
308 /// Groups elements from the `GroupingMap` source by key and finds the maximum of each group
309 /// with respect to the specified comparison function.
310 ///
311 /// If several elements are equally maximum, the last element is picked.
312 ///
313 /// Returns a `HashMap` associating the key of each group with the maximum of that group's elements.
314 ///
315 /// ```
316 /// use itertools::Itertools;
317 ///
318 /// let lookup = vec![1, 3, 4, 5, 7, 8, 9, 12].into_iter()
319 /// .into_grouping_map_by(|&n| n % 3)
320 /// .max_by(|_key, x, y| y.cmp(x));
321 ///
322 /// assert_eq!(lookup[&0], 3);
323 /// assert_eq!(lookup[&1], 1);
324 /// assert_eq!(lookup[&2], 5);
325 /// assert_eq!(lookup.len(), 3);
326 /// ```
327 pub fn max_by<F>(self, mut compare: F) -> HashMap<K, V>
328 where
329 F: FnMut(&K, &V, &V) -> Ordering,
330 {
331 self.reduce(|acc, key, val| match compare(key, &acc, &val) {
332 Ordering::Less | Ordering::Equal => val,
333 Ordering::Greater => acc,
334 })
335 }
336
337 /// Groups elements from the `GroupingMap` source by key and finds the element of each group
338 /// that gives the maximum from the specified function.
339 ///
340 /// If several elements are equally maximum, the last element is picked.
341 ///
342 /// Returns a `HashMap` associating the key of each group with the maximum of that group's elements.
343 ///
344 /// ```
345 /// use itertools::Itertools;
346 ///
347 /// let lookup = vec![1, 3, 4, 5, 7, 8, 9, 12].into_iter()
348 /// .into_grouping_map_by(|&n| n % 3)
349 /// .max_by_key(|_key, &val| val % 4);
350 ///
351 /// assert_eq!(lookup[&0], 3);
352 /// assert_eq!(lookup[&1], 7);
353 /// assert_eq!(lookup[&2], 5);
354 /// assert_eq!(lookup.len(), 3);
355 /// ```
356 pub fn max_by_key<F, CK>(self, mut f: F) -> HashMap<K, V>
357 where
358 F: FnMut(&K, &V) -> CK,
359 CK: Ord,
360 {
361 self.max_by(|key, v1, v2| f(key, v1).cmp(&f(key, v2)))
362 }
363
364 /// Groups elements from the `GroupingMap` source by key and finds the minimum of each group.
365 ///
366 /// If several elements are equally minimum, the first element is picked.
367 ///
368 /// Returns a `HashMap` associating the key of each group with the minimum of that group's elements.
369 ///
370 /// ```
371 /// use itertools::Itertools;
372 ///
373 /// let lookup = vec![1, 3, 4, 5, 7, 8, 9, 12].into_iter()
374 /// .into_grouping_map_by(|&n| n % 3)
375 /// .min();
376 ///
377 /// assert_eq!(lookup[&0], 3);
378 /// assert_eq!(lookup[&1], 1);
379 /// assert_eq!(lookup[&2], 5);
380 /// assert_eq!(lookup.len(), 3);
381 /// ```
382 pub fn min(self) -> HashMap<K, V>
383 where
384 V: Ord,
385 {
386 self.min_by(|_, v1, v2| V::cmp(v1, v2))
387 }
388
389 /// Groups elements from the `GroupingMap` source by key and finds the minimum of each group
390 /// with respect to the specified comparison function.
391 ///
392 /// If several elements are equally minimum, the first element is picked.
393 ///
394 /// Returns a `HashMap` associating the key of each group with the minimum of that group's elements.
395 ///
396 /// ```
397 /// use itertools::Itertools;
398 ///
399 /// let lookup = vec![1, 3, 4, 5, 7, 8, 9, 12].into_iter()
400 /// .into_grouping_map_by(|&n| n % 3)
401 /// .min_by(|_key, x, y| y.cmp(x));
402 ///
403 /// assert_eq!(lookup[&0], 12);
404 /// assert_eq!(lookup[&1], 7);
405 /// assert_eq!(lookup[&2], 8);
406 /// assert_eq!(lookup.len(), 3);
407 /// ```
408 pub fn min_by<F>(self, mut compare: F) -> HashMap<K, V>
409 where
410 F: FnMut(&K, &V, &V) -> Ordering,
411 {
412 self.reduce(|acc, key, val| match compare(key, &acc, &val) {
413 Ordering::Less | Ordering::Equal => acc,
414 Ordering::Greater => val,
415 })
416 }
417
418 /// Groups elements from the `GroupingMap` source by key and finds the element of each group
419 /// that gives the minimum from the specified function.
420 ///
421 /// If several elements are equally minimum, the first element is picked.
422 ///
423 /// Returns a `HashMap` associating the key of each group with the minimum of that group's elements.
424 ///
425 /// ```
426 /// use itertools::Itertools;
427 ///
428 /// let lookup = vec![1, 3, 4, 5, 7, 8, 9, 12].into_iter()
429 /// .into_grouping_map_by(|&n| n % 3)
430 /// .min_by_key(|_key, &val| val % 4);
431 ///
432 /// assert_eq!(lookup[&0], 12);
433 /// assert_eq!(lookup[&1], 4);
434 /// assert_eq!(lookup[&2], 8);
435 /// assert_eq!(lookup.len(), 3);
436 /// ```
437 pub fn min_by_key<F, CK>(self, mut f: F) -> HashMap<K, V>
438 where
439 F: FnMut(&K, &V) -> CK,
440 CK: Ord,
441 {
442 self.min_by(|key, v1, v2| f(key, v1).cmp(&f(key, v2)))
443 }
444
445 /// Groups elements from the `GroupingMap` source by key and find the maximum and minimum of
446 /// each group.
447 ///
448 /// If several elements are equally maximum, the last element is picked.
449 /// If several elements are equally minimum, the first element is picked.
450 ///
451 /// See [`Itertools::minmax`](crate::Itertools::minmax) for the non-grouping version.
452 ///
453 /// Differences from the non grouping version:
454 /// - It never produces a `MinMaxResult::NoElements`
455 /// - It doesn't have any speedup
456 ///
457 /// Returns a `HashMap` associating the key of each group with the minimum and maximum of that group's elements.
458 ///
459 /// ```
460 /// use itertools::Itertools;
461 /// use itertools::MinMaxResult::{OneElement, MinMax};
462 ///
463 /// let lookup = vec![1, 3, 4, 5, 7, 9, 12].into_iter()
464 /// .into_grouping_map_by(|&n| n % 3)
465 /// .minmax();
466 ///
467 /// assert_eq!(lookup[&0], MinMax(3, 12));
468 /// assert_eq!(lookup[&1], MinMax(1, 7));
469 /// assert_eq!(lookup[&2], OneElement(5));
470 /// assert_eq!(lookup.len(), 3);
471 /// ```
472 pub fn minmax(self) -> HashMap<K, MinMaxResult<V>>
473 where
474 V: Ord,
475 {
476 self.minmax_by(|_, v1, v2| V::cmp(v1, v2))
477 }
478
479 /// Groups elements from the `GroupingMap` source by key and find the maximum and minimum of
480 /// each group with respect to the specified comparison function.
481 ///
482 /// If several elements are equally maximum, the last element is picked.
483 /// If several elements are equally minimum, the first element is picked.
484 ///
485 /// It has the same differences from the non-grouping version as `minmax`.
486 ///
487 /// Returns a `HashMap` associating the key of each group with the minimum and maximum of that group's elements.
488 ///
489 /// ```
490 /// use itertools::Itertools;
491 /// use itertools::MinMaxResult::{OneElement, MinMax};
492 ///
493 /// let lookup = vec![1, 3, 4, 5, 7, 9, 12].into_iter()
494 /// .into_grouping_map_by(|&n| n % 3)
495 /// .minmax_by(|_key, x, y| y.cmp(x));
496 ///
497 /// assert_eq!(lookup[&0], MinMax(12, 3));
498 /// assert_eq!(lookup[&1], MinMax(7, 1));
499 /// assert_eq!(lookup[&2], OneElement(5));
500 /// assert_eq!(lookup.len(), 3);
501 /// ```
502 pub fn minmax_by<F>(self, mut compare: F) -> HashMap<K, MinMaxResult<V>>
503 where
504 F: FnMut(&K, &V, &V) -> Ordering,
505 {
506 self.aggregate(|acc, key, val| {
507 Some(match acc {
508 Some(MinMaxResult::OneElement(e)) => {
509 if compare(key, &val, &e) == Ordering::Less {
510 MinMaxResult::MinMax(val, e)
511 } else {
512 MinMaxResult::MinMax(e, val)
513 }
514 }
515 Some(MinMaxResult::MinMax(min, max)) => {
516 if compare(key, &val, &min) == Ordering::Less {
517 MinMaxResult::MinMax(val, max)
518 } else if compare(key, &val, &max) != Ordering::Less {
519 MinMaxResult::MinMax(min, val)
520 } else {
521 MinMaxResult::MinMax(min, max)
522 }
523 }
524 None => MinMaxResult::OneElement(val),
525 Some(MinMaxResult::NoElements) => unreachable!(),
526 })
527 })
528 }
529
530 /// Groups elements from the `GroupingMap` source by key and find the elements of each group
531 /// that gives the minimum and maximum from the specified function.
532 ///
533 /// If several elements are equally maximum, the last element is picked.
534 /// If several elements are equally minimum, the first element is picked.
535 ///
536 /// It has the same differences from the non-grouping version as `minmax`.
537 ///
538 /// Returns a `HashMap` associating the key of each group with the minimum and maximum of that group's elements.
539 ///
540 /// ```
541 /// use itertools::Itertools;
542 /// use itertools::MinMaxResult::{OneElement, MinMax};
543 ///
544 /// let lookup = vec![1, 3, 4, 5, 7, 9, 12].into_iter()
545 /// .into_grouping_map_by(|&n| n % 3)
546 /// .minmax_by_key(|_key, &val| val % 4);
547 ///
548 /// assert_eq!(lookup[&0], MinMax(12, 3));
549 /// assert_eq!(lookup[&1], MinMax(4, 7));
550 /// assert_eq!(lookup[&2], OneElement(5));
551 /// assert_eq!(lookup.len(), 3);
552 /// ```
553 pub fn minmax_by_key<F, CK>(self, mut f: F) -> HashMap<K, MinMaxResult<V>>
554 where
555 F: FnMut(&K, &V) -> CK,
556 CK: Ord,
557 {
558 self.minmax_by(|key, v1, v2| f(key, v1).cmp(&f(key, v2)))
559 }
560
561 /// Groups elements from the `GroupingMap` source by key and sums them.
562 ///
563 /// This is just a shorthand for `self.reduce(|acc, _, val| acc + val)`.
564 /// It is more limited than `Iterator::sum` since it doesn't use the `Sum` trait.
565 ///
566 /// Returns a `HashMap` associating the key of each group with the sum of that group's elements.
567 ///
568 /// ```
569 /// use itertools::Itertools;
570 ///
571 /// let lookup = vec![1, 3, 4, 5, 7, 8, 9, 12].into_iter()
572 /// .into_grouping_map_by(|&n| n % 3)
573 /// .sum();
574 ///
575 /// assert_eq!(lookup[&0], 3 + 9 + 12);
576 /// assert_eq!(lookup[&1], 1 + 4 + 7);
577 /// assert_eq!(lookup[&2], 5 + 8);
578 /// assert_eq!(lookup.len(), 3);
579 /// ```
580 pub fn sum(self) -> HashMap<K, V>
581 where
582 V: Add<V, Output = V>,
583 {
584 self.reduce(|acc, _, val| acc + val)
585 }
586
587 /// Groups elements from the `GroupingMap` source by key and multiply them.
588 ///
589 /// This is just a shorthand for `self.reduce(|acc, _, val| acc * val)`.
590 /// It is more limited than `Iterator::product` since it doesn't use the `Product` trait.
591 ///
592 /// Returns a `HashMap` associating the key of each group with the product of that group's elements.
593 ///
594 /// ```
595 /// use itertools::Itertools;
596 ///
597 /// let lookup = vec![1, 3, 4, 5, 7, 8, 9, 12].into_iter()
598 /// .into_grouping_map_by(|&n| n % 3)
599 /// .product();
600 ///
601 /// assert_eq!(lookup[&0], 3 * 9 * 12);
602 /// assert_eq!(lookup[&1], 1 * 4 * 7);
603 /// assert_eq!(lookup[&2], 5 * 8);
604 /// assert_eq!(lookup.len(), 3);
605 /// ```
606 pub fn product(self) -> HashMap<K, V>
607 where
608 V: Mul<V, Output = V>,
609 {
610 self.reduce(|acc, _, val| acc * val)
611 }
612}
613