Optimize 2024 day 5

The provided rules form a total order for each of the provided update lists, allowing use of the standard library sort functions

Author: ictrobot

crates/year2024/src/day05.rs

@@ -1,46 +1,48 @@
+use std::cmp::Ordering;
 use utils::prelude::*;
 
-/// Sorting lists using a partial order of constraints.
+/// Sorting lists using constraints.
+///
+/// The solution assumes that the rules form a total order for the elements in each sequence.
 #[derive(Clone, Debug)]
 pub struct Day05 {
-    rules: Rules,
+    before: Rules,
     sorted: Vec<Vec<u32>>,
     unsorted: Vec<Vec<u32>>,
 }
 
 const MIN_NUM: usize = 10;
 const MAX_NUM: usize = 99;
 const RANGE: usize = MAX_NUM - MIN_NUM + 1;
-type Rules = [bool; RANGE * RANGE];
+type Rules = [[bool; RANGE]; RANGE];
 
 impl Day05 {
     pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
         let num = parser::number_range(MIN_NUM as u32..=MAX_NUM as u32);
         let rules_parser = num.then(num.with_prefix(b'|')).repeat(b'\n', 1);
         let updates_parser = num.repeat(b',', 1).repeat(b'\n', 1);
 
-        let (rules_list, updates) = rules_parser
+        let (rule_list, updates) = rules_parser
             .then(updates_parser.with_prefix(b'\n'))
             .parse_complete(input)?;
 
-        let mut rules = [false; RANGE * RANGE];
-        for (a, b) in rules_list {
-            rules[(a as usize - MIN_NUM) * RANGE + b as usize] = true;
+        let mut before: Rules = [[false; RANGE]; RANGE];
+        for (a, b) in rule_list {
+            if before[a as usize - MIN_NUM][b as usize - MIN_NUM] {
+                return Err(InputError::new(input, 0, "duplicate rule"));
+            } else if before[b as usize - MIN_NUM][a as usize - MIN_NUM] {
+                return Err(InputError::new(input, 0, "contradictory pair of rules"));
+            }
+
+            before[a as usize - MIN_NUM][b as usize - MIN_NUM] = true;
         }
 
         let (sorted, unsorted) = updates.into_iter().partition(|update| {
-            for (i, &page1) in update.iter().enumerate() {
-                for &page2 in update.iter().skip(i + 1) {
-                    if Self::must_be_before(&rules, page2, page1) {
-                        return false;
-                    }
-                }
-            }
-            true
+            update.is_sorted_by(|&a, &b| before[a as usize - MIN_NUM][b as usize - MIN_NUM])
         });
 
         Ok(Self {
-            rules,
+            before,
             sorted,
             unsorted,
         })
@@ -56,38 +58,24 @@ impl Day05 {
 
     #[must_use]
     pub fn part2(&self) -> u32 {
-        let mut result = 0;
-        let mut list = Vec::new();
-        for original in &self.unsorted {
-            original.clone_into(&mut list);
-
-            for _ in 0..(list.len() / 2) {
-                let i = Self::find_first_index(&self.rules, &list);
-                list.swap_remove(i);
-            }
-
-            let i = Self::find_first_index(&self.rules, &list);
-            result += list[i];
-        }
-
-        result
-    }
-
-    fn must_be_before(rules: &Rules, page1: u32, page2: u32) -> bool {
-        rules[(page1 as usize - MIN_NUM) * RANGE + page2 as usize]
-    }
-
-    fn find_first_index(rules: &Rules, list: &[u32]) -> usize {
-        'outer: for (i, &page1) in list.iter().enumerate() {
-            for (j, &page2) in list.iter().enumerate() {
-                if i != j && Self::must_be_before(rules, page2, page1) {
-                    continue 'outer;
-                }
-            }
-
-            return i;
-        }
-        panic!("no solution found");
+        self.unsorted
+            .iter()
+            .cloned()
+            .map(|mut update| {
+                let index = update.len() / 2;
+                // Will panic if the provided rules are not a total order
+                let (_, middle, _) = update.select_nth_unstable_by(index, |&a, &b| {
+                    if a == b {
+                        Ordering::Equal
+                    } else if self.before[a as usize - MIN_NUM][b as usize - MIN_NUM] {
+                        Ordering::Less
+                    } else {
+                        Ordering::Greater
+                    }
+                });
+                *middle
+            })
+            .sum()
     }
 }