day14 part2

src/day14.rs

@@ -1,6 +1,8 @@
+use std::collections::HashMap;
+
 use aoc_runner_derive::{aoc, aoc_generator};
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
 enum Rock {
     Round,
     Square,
@@ -43,10 +45,8 @@ fn part1(input: &Vec<Vec<Rock>>) -> usize {
                 Rock::Round => {
                     cur_weight -= 1;
                     load += cur_weight;
-                    //println!("{}", cur_weight);
                 },
                 Rock::Square => {
-                    //println!("{}", input[0].len() - i);
                     cur_weight = input[0].len() - i;
                 },
                 Rock::None => continue
@@ -56,10 +56,80 @@ fn part1(input: &Vec<Vec<Rock>>) -> usize {
     load
 }
 
-// #[aoc(day14, part2)]
-// fn part2(input: &Vec<Vec<Rock>>) -> String {
-//     todo!()
-// }
+#[aoc(day14, part2)]
+fn part2(input: &Vec<Vec<Rock>>) -> usize {
+    let mut seen_at = HashMap::new();
+    let mut next_y = Vec::<usize>::new();
+    let mut map = input.to_vec();
+
+    // Functions to get a Rock using a rotated coordinate space
+    let cycle_parts: &[(Box<dyn Fn(&mut Vec<Vec<Rock>>, usize, usize) -> &mut Rock>, _, _); 4] = &[
+        (Box::new(|map: &mut Vec<Vec<Rock>>, x: usize, y: usize| &mut map[y][x]),
+         map[0].len(), map.len()),
+        (Box::new(|map: &mut Vec<Vec<Rock>>, x: usize, y: usize| &mut map[x][y]),
+         map.len(), map[0].len()),
+        (Box::new(|map: &mut Vec<Vec<Rock>>, x: usize, y: usize| { let h = map.len(); &mut map[h - 1 - y][x] }),
+         map[0].len(), map.len()),
+        (Box::new(|map: &mut Vec<Vec<Rock>>, x: usize, y: usize| { let w = map[0].len(); &mut map[x][w - 1 - y] }),
+         map.len(), map[0].len()),
+    ];
+
+    let mut cycle = 0;
+    const END: u32 = 1000000000;
+    while cycle < END {
+        // Handle tilts in each direction
+        for (getter, width, height) in cycle_parts {
+            next_y.clear();
+            next_y.resize(*width, 0);
+            for y in 0..*height {
+                for x in 0..*width {
+                    let item = getter(&mut map, x, y);
+                    match *item {
+                        Rock::None => {}
+                        Rock::Square => {
+                            next_y[x] = y + 1;
+                        }
+                        Rock::Round => {
+                            *item = Rock::None;
+                            *getter(&mut map, x, next_y[x]) = Rock::Round;
+                            next_y[x] += 1;
+                        }
+                    }
+                }
+            }
+        }
+
+        // More compact representation of the current state, for saving in hashmap
+        let key = map.iter().enumerate().flat_map(|(y, line)| {
+            line.iter().enumerate().filter_map(move |(x, &ref rock)| {
+                if rock == &Rock::Round {
+                    Some((x as u8, y as u8))
+                } else {
+                    None
+                }
+            })
+        }).collect::<Vec<_>>();
+
+        cycle += 1;
+        if let Some(seen_at_cycle) = seen_at.insert(key, cycle) {
+            // Current state was identical to one we'd already seen, we can skip forward
+            let diff = cycle - seen_at_cycle;
+            let remaining = END - cycle;
+            let skipped = remaining / diff * diff;
+            cycle += skipped;
+        }
+    }
+
+    let height = map.len();
+
+    map.into_iter().enumerate().flat_map(|(y, line)| line.into_iter().filter_map(move |rock| {
+        if rock == Rock::Round {
+            Some(height - y)
+        } else {
+            None
+        }
+    })).sum()
+}
 
 
 #[cfg(test)]
@@ -82,8 +152,8 @@ O.#..O.#.#
         assert_eq!(part1(&parse(EX)), 136);
     }
 
-    // #[test]
-    // fn part2_example() {
-    //     assert_eq!(part2(&parse(EX)), "<RESULT>");
-    // }
+    #[test]
+    fn part2_example() {
+        assert_eq!(part2(&parse(EX)), 64);
+    }
 }