Rust/2024/01: improve solution

Rust/2024/01.rs

@@ -1,5 +1,6 @@
 #![feature(test)]
 
+use aoc::iter_ext::IterExt;
 use itertools::Itertools;
 
 type Input = Vec<(i32, i32)>;
@@ -28,9 +29,10 @@ fn part1(input: &Input) -> u32 {
 }
 
 fn part2(input: &Input) -> i32 {
+    let right_counts = input.iter().map(|&(_, r)| r).counts_fx();
     input
         .iter()
-        .map(|&(l, _)| l * input.iter().filter(|&&(_, r)| l == r).count() as i32)
+        .map(|&(l, _)| l * *right_counts.get(&l).unwrap_or(&0) as i32)
         .sum()
 }
 

Rust/lib/iter_ext.rs

@@ -1,4 +1,6 @@
-use std::iter::Iterator;
+use std::{hash::Hash, iter::Iterator};
+
+use rustc_hash::FxHashMap;
 
 pub trait IterExt: Iterator {
     fn take_while_inclusive<P>(self, predicate: P) -> TakeWhileInclusive<Self, P>
@@ -32,6 +34,17 @@ pub trait IterExt: Iterator {
     {
         self.position(|x| elem == x)
     }
+
+    fn counts_fx(self) -> FxHashMap<Self::Item, usize>
+    where
+        Self: Sized,
+        Self::Item: Hash + Eq,
+    {
+        self.fold(FxHashMap::default(), |mut counts, x| {
+            counts.entry(x).and_modify(|x| *x += 1).or_insert(1);
+            counts
+        })
+    }
 }
 
 impl<I> IterExt for I