Dag 9 aoc, zonder trees maar gewoon met een vieze O(n^3) vector

advent_of_code/2024/9/src/main.rs

@@ -1,44 +1,129 @@
-use std::{cell::RefCell, rc::Rc};
+use std::collections::HashSet;
 
 #[derive(Debug, Clone)]
-struct FileTree {
-    val: u32,
-    left: Option<TreeNodeRef>,
-    right: Option<TreeNodeRef>,
+enum File {
+    Empty,
+    File(u64),
 }
 
-type TreeNodeRef = Rc<RefCell<FileTree>>;
-
-impl FileTree {
-    fn new() -> Self {
-        FileTree {
-            val: 0,
-            left: None,
-            right: None,
-        }
-    }
-
-    fn fill(&mut self, filesystem: &[u32]) {
-        for (i, file) in filesystem.iter().enumerate() {
-        }
-    }
-
-    fn add_file(&mut self, length: u32) {}
-    fn add_empty(&mut self, length: u32) {}
+#[derive(Debug, Clone)]
+enum Fs {
+    Empty(u64),
+    File(u64, u64),
 }
 
 fn parse(input: &str) -> Vec<u32> {
-    input.chars().map(|c| c.to_digit(10).unwrap()).collect()
+    input.chars().filter_map(|c| c.to_digit(10)).collect()
 }
 
-fn solve_1(input: &str) -> u32 {
-    let filesystem = parse(input);
-    let mut filetree = FileTree::new();
+fn solve_1(input: &str) -> u64 {
+    let original_filesystem: Vec<_> = parse(input).into_iter().map(|f| f as u64).collect();
 
-    filetree.fill(&filesystem[..]);
+    let mut filesystem = vec![];
+    for (i, d) in original_filesystem.into_iter().enumerate() {
+        if i % 2 == 0 {
+            for _ in 0..d {
+                filesystem.push(File::File(i as u64 / 2));
+            }
+        } else {
+            for _ in 0..d {
+                filesystem.push(File::Empty);
+            }
+        }
+    }
+
+    for ei in 0..filesystem.len() {
+        if let Some(File::Empty) = filesystem.get(ei) {
+            while let Some(File::Empty) = filesystem.last() {
+                filesystem.pop();
+            }
+
+            if filesystem.len() > ei {
+                if let Some(File::File(fid)) = filesystem.pop() {
+                    filesystem[ei] = File::File(fid);
+                }
+            } else {
+                break;
+            }
+        }
+    }
+
+    filesystem.into_iter().enumerate().fold(0, |s, (i, f)| {
+        s + i as u64
+            * match f {
+                File::Empty => 0,
+                File::File(fid) => fid,
+            }
+    })
 }
 
-fn solve_2(input: &str) {}
+fn solve_2(input: &str) -> u64 {
+    let original_filesystem: Vec<_> = parse(input).into_iter().map(|f| f as u64).collect();
+
+    let mut filesystem = vec![];
+    let mut c = 0;
+    for (i, d) in original_filesystem.into_iter().enumerate() {
+        if i % 2 == 0 {
+            assert_ne!(0, d);
+
+            filesystem.push(Fs::File(i as u64 / 2, d));
+            c += 1;
+        } else if d != 0 {
+            filesystem.push(Fs::Empty(d));
+        }
+    }
+
+    for fid in (0..=c).rev() {
+        for i in (0..filesystem.len()).rev() {
+            if let Fs::File(file_id, wf) = filesystem[i] {
+                if file_id == fid {
+                    for j in 0..=i {
+                        if let Fs::Empty(we) = filesystem[j] {
+                            if we >= wf {
+                                filesystem[j] = Fs::File(fid, wf);
+
+                                if i > 0 {
+                                    if let Fs::Empty(el) = filesystem[i - 1] {
+                                        if let Some(Fs::Empty(er)) = filesystem.get(i + 1) {
+                                            filesystem[i - 1] = Fs::Empty(el + wf + er);
+                                            filesystem.remove(i + 1);
+                                            filesystem.remove(i);
+                                        } else {
+                                            filesystem[i - 1] = Fs::Empty(el + wf);
+                                            filesystem.remove(i);
+                                        }
+                                    } else if let Some(Fs::Empty(er)) = filesystem.get(i + 1) {
+                                        filesystem[i] = Fs::Empty(wf + er);
+                                        filesystem.remove(i + 1);
+                                    } else {
+                                        filesystem[i] = Fs::Empty(wf);
+                                    }
+                                }
+
+                                if we > wf {
+                                    filesystem.insert(j + 1, Fs::Empty(we - wf));
+                                }
+                                break;
+                            }
+                        }
+                    }
+                    break;
+                }
+            }
+        }
+    }
+
+    filesystem
+        .into_iter()
+        .fold((0, 0), |(s, i), f| match f {
+            Fs::Empty(w) => (s, i + w as usize),
+            Fs::File(fid, w) => (
+                s + (i..(i + w as usize)).map(|it| it as u64 * fid).sum::<u64>(),
+                i + w as usize,
+            ),
+        })
+        .0
+}
 
 fn main() {
     println!("Hello, this is Patrick!");
@@ -49,10 +134,10 @@ fn main() {
     println!("The checksum is {}", result_1);
 
     let result_2 = solve_2(input);
+    println!("The new checksum is {}", result_2);
 }
 
 #[cfg(test)]
-
 mod tests {
     use super::*;
 
@@ -63,5 +148,8 @@ mod tests {
     }
 
     #[test]
-    fn test_2() {}
+    fn test_2() {
+        let test_input = include_str!("../test.txt");
+        assert_eq!(solve_2(test_input), 2858);
+    }
 }