1 files changed, 23 insertions, 24 deletions

diff --git a/src/bin/day20.rs b/src/bin/day20.rs
index 79ab0f9..48b58b3 100644
--- a/src/bin/day20.rs
+++ b/src/bin/day20.rs
@@ -2,41 +2,39 @@ static DAY: u8 = 20;
 
 fn main() {
     let input = advent::read_lines(DAY);
-    println!("{DAY}a: {}", grove_coordinates(&input));
-    println!("{DAY}b: {}", 0);
+    println!("{DAY}a: {}", grove_coordinates(&input, false));
+    println!("{DAY}b: {}", grove_coordinates(&input, true));
 }
 
-fn grove_coordinates(input: &[String]) -> isize {
+fn grove_coordinates(input: &[String], full_decryption: bool) -> isize {
+    let (decryption_key, rounds) = if full_decryption { (811589153, 10) } else { (1, 1) };
+
     let mut input = input.iter()
-                         .map(|x| (x.parse::<isize>().unwrap(), false))
+                         .enumerate()
+                         .map(|(idx,x)| (idx, x.parse::<isize>().unwrap() * decryption_key))
                          .collect::<Vec<_>>();
 
-    loop {
-        let pos = match input.iter().position(|x| !x.1) {
-            None => break,
-            Some(p) => p,
-        };
-        let val = input[pos];
-        input.remove(pos);
+    for _ in 0 .. rounds {
+        for idx in 0 .. input.len() {
+            let pos = input.iter().position(|x| x.0 == idx).unwrap();
+            let val = input[pos].1;
+            input.remove(pos);
 
-        let input_len = input.len() as isize;
+            let input_len = input.len() as isize;
 
-        let pos_diff = val.0 % input_len;
-        let mut new_pos = (pos as isize + pos_diff) % input_len;
-        if new_pos.is_negative() {
-            new_pos += input_len;
-        }
+            let pos_diff = val % input_len;
+            let mut new_pos = (pos as isize + pos_diff) % input_len;
+            if new_pos <= 0 {
+                new_pos += input_len;
+            }
 
-        if new_pos == 0 {
-            input.push((val.0 as isize, true));
-        } else {
-            input.insert(new_pos as usize, (val.0 as isize, true));
+            input.insert(new_pos as usize, (idx, val));
         }
     }
 
-    let pos = input.iter().position(|x| x.0 == 0).unwrap();
+    let pos = input.iter().position(|x| x.1 == 0).unwrap();
     let n = ((pos + 1000) % input.len(), (pos + 2000) % input.len(), (pos + 3000) % input.len());
-    input[n.0].0 + input[n.1].0 + input[n.2].0
+    input[n.0].1 + input[n.1].1 + input[n.2].1
 }
 
 #[cfg(test)]
@@ -55,6 +53,7 @@ mod tests {
             "4",
         ].iter().map(|&x| String::from(x)).collect::<Vec<_>>();
 
-        assert_eq!(grove_coordinates(&input), 3);
+        assert_eq!(grove_coordinates(&input, false), 3);
+        assert_eq!(grove_coordinates(&input, true), 1623178306);
     }
 }