1 files changed, 36 insertions, 5 deletions

diff --git a/src/bin/day5.rs b/src/bin/day5.rs
index d13be85..e26b1a1 100644
--- a/src/bin/day5.rs
+++ b/src/bin/day5.rs
@@ -5,7 +5,7 @@ static DAY: u8 = 5;
 fn main() {
     let input = advent::read_file(DAY);
     println!("{DAY}a: {}", PageOrdering::new(&input).right_order_middle());
-    println!("{DAY}b: {}", 0);
+    println!("{DAY}b: {}", PageOrdering::new(&input).incorrect_fixed_middle());
 }
 
 struct PageOrdering {
@@ -34,7 +34,7 @@ impl PageOrdering {
         PageOrdering { rules, updates }
     }
 
-    fn correct_order(&self, update: &[u32]) -> bool {
+    fn wrong_position(&self, update: &[u32]) -> Option<u32> {
         for (i, page1) in update.iter().enumerate() {
             for page2 in update.iter().skip(i+1) {
                 let before2 = match self.rules.get(page2) {
@@ -42,22 +42,52 @@ impl PageOrdering {
                     None => continue,
                 };
                 if before2.contains(page1) {
-                    return false;
+                    return Some(*page1);
                 }
             }
         }
-        true
+        None
     }
 
     fn right_order_middle(&self) -> u32 {
         let mut sum = 0;
         for update in &self.updates {
-            if self.correct_order(update) {
+            if self.wrong_position(update).is_none() {
                 sum += update[update.len()/2];
             }
         }
         sum
     }
+
+    fn fix_order(&self, update: &[u32]) -> Vec<u32> {
+        let mut update = update.to_vec();
+        let page_cmp = |page1: u32, page2: u32| -> std::cmp::Ordering {
+            if let Some(rules) = self.rules.get(&page1) {
+                if rules.contains(&page2) {
+                    return std::cmp::Ordering::Less;
+                }
+            }
+            if let Some(rules) = self.rules.get(&page2) {
+                if rules.contains(&page1) {
+                    return std::cmp::Ordering::Greater;
+                }
+            }
+            panic!("no ordering found");
+        };
+        update.sort_by(|page1, page2| page_cmp(*page1, *page2));
+        update
+    }
+
+    fn incorrect_fixed_middle(&self) -> u32 {
+        let mut sum = 0;
+        for update in &self.updates {
+            if self.wrong_position(update).is_some() {
+                let fixed = self.fix_order(update);
+                sum += fixed[fixed.len()/2];
+            }
+        }
+        sum
+    }
 }
 
 #[cfg(test)]
@@ -96,5 +126,6 @@ mod tests {
 61,13,29
 97,13,75,29,47";
         assert_eq!(PageOrdering::new(input).right_order_middle(), 143);
+        assert_eq!(PageOrdering::new(input).incorrect_fixed_middle(), 123);
     }
 }