Optimize & simplify day 3

bram/rust/src/day_03/mod.rs

@@ -1,66 +1,43 @@
 use crate::utils;
 
-pub fn answer(text : String) -> ( u16, u64 ) {
+pub fn answer(text : String) -> ( u64, u64 ) {
-    ( text.as_str().split_whitespace().map(BatteryBank::from).sum()
+    text.as_str()
-    , text.as_str().split_whitespace().map(BigBatteryBank::from).sum()
+        .split_whitespace()
-    )
+        .map(|s|
+            ( BatteryBank::from(s, 2), BatteryBank::from(s, 12) )
+        )
+        .fold((0, 0), |(t1, t2), (b1, b2)| (t1 + b1, t2 + b2))
 }
 
 struct BatteryBank {
-    highest_digit : u8,
+    digits : Vec<u8>,
-    highest_value : u8,
+    size : usize,
 }
 impl BatteryBank {
-    fn add_battery(&mut self, b : u8) {
-        let joltage = 10 * self.highest_digit + b;
-
-        if joltage > self.highest_value {
-            self.highest_value = joltage;
-        }
-        if b > self.highest_digit {
-            self.highest_digit = b;
-        }
-    }
-
-    fn from(s : &str) -> u16 {
-        let mut bank = BatteryBank{ highest_digit : 0, highest_value : 0 };
-
-        for n in s.chars().filter_map(utils::char_to_u8) {
-            bank.add_battery(n);
-        }
-
-        bank.highest_value as u16
-    }
-}
-
-struct BigBatteryBank {
-    highest_digit : Vec<u8>,
-}
-impl BigBatteryBank {
     fn add_battery(&mut self, b : u8, min_pos : usize) {
-        for i in min_pos..12 {
+        for i in min_pos..self.size {
-            match self.highest_digit.get(i) {
+            match self.digits.get(i) {
                 Some(x) =>
                     if b > *x {
-                        self.highest_digit.truncate(i);
+                        self.digits.truncate(i);
-                        self.highest_digit.push(b);
+                        self.digits.push(b);
                         return;
                     },
                 None => {
-                    self.highest_digit.push(b);
+                    self.digits.push(b);
                     return;
                 },
             }
         }
     }
 
-    fn from(s : &str) -> u64 {
+    fn from(s : &str, size : usize) -> u64 {
-        let mut bank = BigBatteryBank { highest_digit : Vec::new() };
+        let mut bank = BatteryBank { digits : Vec::new(), size : size };
         let l : usize = s.len();
 
         for (i, n) in s.chars().filter_map(utils::char_to_u8).enumerate() {
             let offset : usize = l - i;
-            bank.add_battery(n, if offset > 12 { 0 } else { 12 - offset });
+            bank.add_battery(n, if offset > size { 0 } else { size - offset });
         }
 
         bank.measure_joltage().unwrap()
@@ -69,8 +46,8 @@ impl BigBatteryBank {
     fn measure_joltage(&self) -> Option<u64> {
         let mut joltage : u64 = 0;
 
-        for i in 0..12 {
+        for i in 0..self.size {
-            let n : &u8 = self.highest_digit.get(i)?;
+            let n : &u8 = self.digits.get(i)?;
             joltage =  10 * joltage + (*n as u64);
         }