refactor 01a and 01b using itertools

5 files changed, 79 insertions, 83 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 8dad0ca..cb5c260 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -69,6 +69,7 @@ name = "day01a"
 version = "0.1.0"
 dependencies = [
  "clap",
+ "itertools",
 ]
 
 [[package]]
@@ -76,6 +77,7 @@ name = "day01b"
 version = "0.1.0"
 dependencies = [
  "clap",
+ "itertools",
 ]
 
 [[package]]
@@ -83,6 +85,12 @@ name = "day02a"
 version = "0.1.0"
 
 [[package]]
+name = "either"
+version = "1.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "90e5c1c8368803113bf0c9584fc495a58b86dc8a29edbf8fe877d21d9507e797"
+
+[[package]]
 name = "hashbrown"
 version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -114,6 +122,15 @@ dependencies = [
 ]
 
 [[package]]
+name = "itertools"
+version = "0.10.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473"
+dependencies = [
+ "either",
+]
+
+[[package]]
 name = "libc"
 version = "0.2.138"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/day01a/Cargo.toml b/day01a/Cargo.toml
index d768066..cc1c83a 100644
--- a/day01a/Cargo.toml
+++ b/day01a/Cargo.toml
@@ -7,3 +7,4 @@ edition = "2021"
 
 [dependencies]
 clap = { version = "3.2.20", features = ["derive"] }
+itertools = "0.10.5"
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()
-    );
 }
diff --git a/day01b/Cargo.toml b/day01b/Cargo.toml
index 4fcdb65..73ea6d0 100644
--- a/day01b/Cargo.toml
+++ b/day01b/Cargo.toml
@@ -7,3 +7,4 @@ edition = "2021"
 
 [dependencies]
 clap = { version = "3.2.20", features = ["derive"] }
+itertools = "0.10.5"
diff --git a/day01b/src/main.rs b/day01b/src/main.rs
index 4c880c4..1f93202 100644
--- a/day01b/src/main.rs
+++ b/day01b/src/main.rs
@@ -16,6 +16,7 @@
 /// Find the top three Elves carrying the most Calories.  How many Calories are those Elves carrying
 /// in total?
 use clap::Parser;
+use itertools::Itertools;
 
 use std::collections::BinaryHeap;
 use std::fs::File;
@@ -54,53 +55,41 @@ 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();
-    }
-
-    // Find the elf carrying the most calories
+    let _ = reader
+        .lines()
+        .map(|l| l.unwrap())
+        .peekable()
+        .batching(|it| {
+            if it.peek().is_none() {
+                return None;
+            }
+            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.iter().sum(),
+            });
+            match heap.peek() {
+                None => None,
+                Some(_) => Some(()),
+            }
+        })
+        .last();
+
+    // Find the top 3 elves carrying the most calories
     let res = heap
         .into_iter_sorted()
         .inspect(|x| println!("calories={:6} elf={}", x.calories, x.elf))
         .take(3)
         .fold(0usize, |acc, x| acc + x.calories);
     println!("The top 3 elves are carry {res} calories");
-
-    let mut cals = include_str!("../examples/input.txt")
-        .split("\n\n")
-        .map(|e| e.lines().map(|c| c.parse::<u32>().unwrap()).sum())
-        .collect::<Vec<u32>>();
-    cals.sort_unstable();
-    println!("{}", cals.into_iter().rev().take(3).sum::<u32>());
 }