rust 2024-05 use nth sorted

1 files changed, 11 insertions, 11 deletions

diff --git a/2024/05/rust/src/main.rs b/2024/05/rust/src/main.rs
index 50c9707..0db788d 100644
--- a/2024/05/rust/src/main.rs
+++ b/2024/05/rust/src/main.rs
@@ -1,16 +1,16 @@
 #![feature(slice_split_once)]
 
-use std::cmp::Ordering;
+use std::cmp::Ordering::{Greater, Less};
 use std::collections::HashSet;
 use std::io;
 
 fn main() -> io::Result<()> {
     let mut lines = io::stdin().lines().flatten();
-    let mut heap = HashSet::<(u16, u16)>::new();
+    let mut set = HashSet::<(u16, u16)>::new();
 
     while let Some(s) = lines.next() {
         if let Some((a, b)) = s.split_once('|') {
-            heap.insert((a.parse().unwrap(), b.parse().unwrap()));
+            set.insert((a.parse().unwrap(), b.parse().unwrap()));
         } else {
             break;
         }
@@ -21,17 +21,17 @@ fn main() -> io::Result<()> {
 
     for line in lines {
         let mut nums: Vec<u16> = line.split(',').map(|n| n.parse().unwrap()).collect();
-        let presort = nums.clone();
         let mid = nums.len() / 2;
-        nums.sort_unstable_by(|&a, &b| match heap.contains(&(a, b)) {
-            true => Ordering::Less,
-            false => Ordering::Greater,
-        });
 
-        if nums.iter().eq(presort.iter()) {
-            silver += presort[mid] as usize;
+        if nums.is_sorted_by(|&a, &b| set.contains(&(a, b))) {
+            silver += nums[mid] as usize;
         } else {
-            gold += nums[mid] as usize;
+            gold += *nums
+                .select_nth_unstable_by(mid, |&a, &b| match set.contains(&(a, b)) {
+                    true => Less,
+                    false => Greater,
+                })
+                .1 as usize;
         }
     }