day 9 but very yikes

Aoc2024/Day09/Main.lean

@@ -16,4 +16,4 @@ def main : IO Unit := do
   IO.println s!"Example: {<- parseAndSolvePart2 exampleInput}"
   let answerPart2 <- parseAndSolvePart2 puzzleInput
   IO.println s!"Puzzle: {answerPart2}"
-  assert! (answerPart2 == -1)
+  assert! (answerPart2 == 6272188244509)

Aoc2024/Day09/Solve.lean

@@ -12,6 +12,7 @@ private def blocksToString (blocks : Blocks) : String :=
     | none => "."
     | some id => toString id
   b.foldl (· ++ ·) ""
+
 #guard blocksToString #[none, none, some 1, none, some 2] == "..1.2"
 
 inductive ExpandMode where
@@ -31,7 +32,7 @@ private def expandDiskMap (diskMap : List Nat) : Blocks :=
 private def nextFreeBlockOnOrAfter (i : Nat) (blocks : Blocks) : Nat :=
   let rec loop (i : Nat) : Nat :=
     if i < blocks.size then
-      match blocks.get! i with
+      match blocks.getD i none with
       | none => i
       | some _ => loop (i + 1)
     else
@@ -43,7 +44,7 @@ private def nextFreeBlockOnOrAfter (i : Nat) (blocks : Blocks) : Nat :=
 private def nextFileBlockOnOrBefore (i : Nat) (blocks : Blocks) : Nat :=
   let rec loop (i : Nat) : Nat :=
     if i > 0 then
-      match blocks.get! i with
+      match blocks.getD i none with
       | none => loop (i - 1)
       | some _ => i
     else
@@ -62,14 +63,14 @@ private def swapBlocks (blocks : Blocks) (index1 index2 : Nat) : Blocks :=
   #[some 0, some 2, none, none, none, some 1]
 
 private def checksum (blocks : Blocks) : Int :=
-  blocks.toList.enum.sumBy fun
+  blocks.toList.enum.sumBy λ
     | (_, none) => 0
     | (i, some id) => id * i
 
 private def solvePart1 (diskMap : List Nat) : Int := Id.run do
   let mut blocks := expandDiskMap diskMap
   let mut nextFree := nextFreeBlockOnOrAfter 0 blocks
-  let mut nextFile := blocks.size - 1 -- assumption on input: ends with file blocks
+  let mut nextFile := nextFileBlockOnOrBefore (blocks.size - 1) blocks
   while nextFree < nextFile do
     blocks := swapBlocks blocks nextFree nextFile
     nextFree := nextFreeBlockOnOrAfter (nextFree + 1) blocks
@@ -82,8 +83,68 @@ def parseAndSolvePart1 (s : String): Except String Int := parseDiskMap s |>.map
 
 -- Part 2
 
-private def solvePart2 (diskMap : List Nat) : Int := sorry
+private def swapBlockRange (blocks : Blocks) (index1 index2 length : Nat) : Blocks := Id.run do
+  let mut blocks := blocks
+  for i in [0:length] do
+    blocks := swapBlocks blocks (index1 + i) (index2 + i)
+  blocks
+
+#guard swapBlockRange #[some 0, none, none, none, some 9, some 9, none] (index1 := 1) (index2 := 4) (length := 2) ==
+  #[some 0, some 9, some 9, none, none, none, none]
+
+private def getHighestId (blocks : Blocks) : IdNum :=
+  blocks.foldl (fun max block => match block with
+    | none => max
+    | some id => Nat.max max id) 0
+
+structure Range where
+  index: Nat
+  length: Nat
+deriving Repr, BEq
+
+private def locateFile (blocks : Blocks) (id : IdNum) : Range :=
+  let index := blocks.findIdx? (λ block => block == some id) |>.getD 0
+  let length: Nat := Id.run do
+    let mut length := 1
+    while blocks.getD (index + length) none == some id do
+      length := length + 1
+    length
+  { index, length }
+
+#guard locateFile #[none, some 0, some 0, none, none, none, some 1] 0 == { index := 1, length := 2 }
+
+private def leftmostFreeSpace (blocks : Blocks) (length : Nat): Option Nat :=
+  let rec loop (n : Nat) : Option Nat :=
+    match n with
+    | 0 => none
+    | n' + 1 =>
+      let i := blocks.size - n
+      if blocks.toSubarray i (i + length) |>.all (· == none) then
+        some i
+      else
+        loop n'
+  loop blocks.size
+
+#guard leftmostFreeSpace #[some 0, none, none, some 1, none, none, none, some 2] 2 == some 1
+#guard leftmostFreeSpace #[some 0, none, none, some 1, none, none, none, some 2] 3 == some 4
+
+private def solvePart2 (diskMap : List Nat) : Int := Id.run do
+  let mut blocks := expandDiskMap diskMap
+  let mut currentId := getHighestId blocks
+  while true do
+    -- dbg_trace blocksToString blocks
+    let range := locateFile blocks currentId
+    let maybeSpace := leftmostFreeSpace blocks range.length
+    match maybeSpace with
+    | none => blocks := blocks
+    | some spaceIndex =>
+      if spaceIndex < range.index then
+        blocks := swapBlockRange blocks spaceIndex range.index range.length
+    if currentId == 0 then
+      break
+    currentId := currentId - 1
+  checksum blocks
 
 def parseAndSolvePart2 (s : String): Except String Int := parseDiskMap s |>.map solvePart2
 
--- #guard parseAndSolvePart2 exampleInput == Except.ok -1
+#guard parseAndSolvePart2 exampleInput == Except.ok 2858