Day 9: Disk Defragmenter

README.md

@@ -13,3 +13,4 @@ Development occurs in language-specific directories:
 |[Day6.hs](hs/src/Day6.hs)|[Day6.kt](kt/aoc2024-lib/src/commonMain/kotlin/com/github/ephemient/aoc2024/Day6.kt)|[day6.py](py/aoc2024/day6.py)|[day6.rs](rs/src/day6.rs)|
 |[Day7.hs](hs/src/Day7.hs)|[Day7.kt](kt/aoc2024-lib/src/commonMain/kotlin/com/github/ephemient/aoc2024/Day7.kt)|[day7.py](py/aoc2024/day7.py)|[day7.rs](rs/src/day7.rs)|
 |[Day8.hs](hs/src/Day8.hs)|[Day8.kt](kt/aoc2024-lib/src/commonMain/kotlin/com/github/ephemient/aoc2024/Day8.kt)|[day8.py](py/aoc2024/day8.py)|[day8.rs](rs/src/day8.rs)|
+|[Day9.hs](hs/src/Day9.hs)||||

hs/aoc2024.cabal

@@ -29,6 +29,7 @@ library
     Day6
     Day7
     Day8
+    Day9
 
   other-modules:
     Common
@@ -38,6 +39,7 @@ library
     containers ^>=0.7,
     megaparsec ^>=9.7.0,
     parallel ^>=3.2.2.0,
+    split ^>=0.2.5,
     text ^>=2.1.2,
     vector ^>=0.13.2.0,
 
@@ -75,6 +77,7 @@ test-suite aoc2024-test
     Day6Spec
     Day7Spec
     Day8Spec
+    Day9Spec
 
   build-depends:
     aoc2024,

hs/app/Main.hs

@@ -15,6 +15,7 @@ import Day5 qualified (part1, part2)
 import Day6 qualified (part1, part2)
 import Day7 qualified (part1, part2)
 import Day8 qualified (part1, part2)
+import Day9 qualified (part1, part2)
 import System.Environment (getArgs, lookupEnv)
 import System.FilePath (combine)
 import Text.Megaparsec (errorBundlePretty)
@@ -46,3 +47,4 @@ main = do
   run 6 print [Day6.part1, Day6.part2]
   run 7 (either fail print) [Day7.part1, Day7.part2]
   run 8 print [Day8.part1, Day8.part2]
+  run 9 print [Day9.part1, Day9.part2]

hs/bench/Main.hs

@@ -14,6 +14,7 @@ import Day5 qualified (part1, part2)
 import Day6 qualified (part1, part2)
 import Day7 qualified (part1, part2)
 import Day8 qualified (part1, part2)
+import Day9 qualified (part1, part2)
 import System.Environment.Blank (getEnv, setEnv, unsetEnv)
 import System.FilePath (combine)
 
@@ -78,5 +79,11 @@ main =
           "Day 8"
           [ bench "part 1" $ nf Day8.part1 input,
             bench "part 2" $ nf Day8.part2 input
+          ],
+      env (getDayInput 9) $ \input ->
+        bgroup
+          "Day 9"
+          [ bench "part 1" $ nf Day9.part1 input,
+            bench "part 2" $ nf Day9.part2 input
           ]
     ]

hs/src/Day9.hs

@@ -0,0 +1,56 @@
+{-# LANGUAGE ParallelListComp #-}
+
+-- |
+-- Module:         Day9
+-- Description:    <https://adventofcode.com/2024/day/9 Day 9: Disk Defragmenter>
+module Day9 (part1, part2) where
+
+import Control.Monad (ap)
+import Data.Char (digitToInt, isDigit)
+import Data.List (mapAccumR, scanl')
+import Data.List.Split (chunksOf)
+import Data.Sequence qualified as Seq (Seq ((:<|), (:|>)), fromList, length, spanl, take, (<|), (|>))
+import Data.Text (Text)
+import Data.Text qualified as T (unpack)
+
+triRange :: (Integral a) => a -> a -> a
+triRange offset size = (2 * offset + size - 1) * size `div` 2
+
+part1 :: Text -> Int
+part1 input = snd $ foldl' add (0, 0) disk1
+  where
+    disk0 = zip ([0 ..] >>= (: [Nothing]) . Just) . map digitToInt . filter isDigit $ T.unpack input
+    disk1 = defrag . Seq.fromList $ disk0
+    defrag (disk Seq.:|> (Nothing, _)) = defrag disk
+    defrag ((Just name, size) Seq.:<| disk) = (name, size) : defrag disk
+    defrag ((Nothing, freeSize) Seq.:<| (disk Seq.:|> (Just name, size)))
+      | size <= freeSize = (name, size) : defrag ((Nothing, freeSize - size) Seq.<| disk)
+      | otherwise = (name, freeSize) : defrag (disk Seq.|> (Just name, size - freeSize))
+    defrag _ = []
+    add (offset, !total) (name, size) = (offset + size, total + name * triRange offset size)
+
+part2 :: Text -> Int
+part2 input =
+  sum
+    [ name * triRange offset size
+    | (name, offset, size) <- snd $ mapAccumR defrag free0 disk0
+    ]
+  where
+    disk0 =
+      [ (name, offset, size)
+      | name <- [0 ..]
+      | (size, offset) : _ <- chunksOf 2 . (zip `ap` scanl' (+) 0) . map digitToInt $ T.unpack input
+      ]
+    free0 =
+      Seq.fromList
+        [ (start, end - start)
+        | (_, offset, size) <- disk0,
+          let start = offset + size
+        | (_, end, _) <- drop 1 disk0
+        ]
+    defrag free (name, _, size)
+      | (free1, (freeOffset, freeSize) Seq.:<| free2) <- Seq.spanl (\(_, freeSize) -> freeSize < size) free =
+          ( free1 <> Seq.take (Seq.length free2) ((freeOffset + size, freeSize - size) Seq.<| free2),
+            (name, freeOffset, size)
+          )
+    defrag free file = (Seq.take (Seq.length free - 1) free, file)

hs/test/Day9Spec.hs

@@ -0,0 +1,23 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module Day9Spec (spec) where
+
+import Data.Text (Text)
+import Data.Text qualified as T (unlines)
+import Day9 (part1, part2)
+import Test.Hspec (Spec, describe, it, shouldBe)
+
+example :: Text
+example =
+  T.unlines
+    [ "2333133121414131402"
+    ]
+
+spec :: Spec
+spec = do
+  describe "part 1" $ do
+    it "examples" $ do
+      part1 example `shouldBe` 1928
+  describe "part 2" $ do
+    it "examples" $ do
+      part2 example `shouldBe` 2858