refactor 01a and 01b using itertools

1 files changed, 30 insertions, 42 deletions

diff --git a/day01a/src/main.rs b/day01a/src/main.rs
index 64b1c7e..6c6faa7 100644
--- a/day01a/src/main.rs
+++ b/day01a/src/main.rs
@@ -58,6 +58,7 @@
 ///
 /// Find the Elf carrying the most Calories.  How many total Calories is that Elf carrying?
 use clap::Parser;
+use itertools::Itertools;
 
 use std::collections::BinaryHeap;
 use std::fs::File;
@@ -96,43 +97,39 @@ fn main() {
     let args = Cli::parse();
 
     let file = File::open(&args.file).unwrap();
-    let mut reader = BufReader::new(file);
+    let reader = BufReader::new(file);
 
     let mut heap = BinaryHeap::new();
-    let mut line = String::new();
-    let mut calorie_count = CalorieCount {
-        elf: 1,
-        calories: 0,
-    };
-    while let Ok(bytes) = reader.read_line(&mut line) {
-        // EOF
-        if bytes == 0 {
-            heap.push(calorie_count);
-            break;
-        }
-
-        // Newlines
-        let tline = line.trim();
-        if tline.is_empty() {
-            // Push onto the binary heap
-            heap.push(calorie_count);
-            calorie_count.calories = 0;
-            calorie_count.elf += 1;
-            continue;
-        }
-
-        // Parse the calories
-        let calories = tline.parse::<usize>().unwrap();
-
-        // Accumulate the calorie count for this elf
-        calorie_count.calories += calories;
-
-        // Clear the line for the next iteration
-        line.clear();
-    }
+    let res = reader
+        .lines()
+        .map(|l| l.unwrap())
+        .peekable()
+        .batching(|it| {
+            if it.peek().is_none() {
+                return None;
+            }
+            // TODO(shivesh): use the std::iter try_collect method when possible
+            Some(
+                it.take_while(|l| !l.is_empty())
+                    .map(|w| w.parse::<usize>().unwrap())
+                    .collect::<Vec<usize>>(),
+            )
+        })
+        .zip(1..)
+        .scan(&mut heap, |heap, (batch, elf)| {
+            heap.push(CalorieCount {
+                elf,
+                calories: batch.into_iter().sum(),
+            });
+            match heap.peek() {
+                None => None,
+                Some(&v) => Some(v),
+            }
+        })
+        .last();
 
     // Find the elf carrying the most calories
-    let res = heap.peek().unwrap();
+    let res = res.unwrap();
     println!(
         "The most amount of calories carried by an elf is {} and is carried by elf {}",
         res.calories, res.elf
@@ -142,13 +139,4 @@ fn main() {
         .inspect(|x| println!("calories={:6} elf={}", x.calories, x.elf))
         .take(10)
         .last();
-
-    println!(
-        "{}",
-        include_str!("../examples/input.txt")
-            .split("\n\n")
-            .map(|e| e.lines().map(|c| c.parse::<u32>().unwrap()).sum::<u32>())
-            .max()
-            .unwrap()
-    );
 }