[Day 09] Part 2: n**2

1 files changed, 41 insertions, 7 deletions

diff --git a/day09/day09.hs b/day09/day09.hs
index ca87b85..c878584 100644
--- a/day09/day09.hs
+++ b/day09/day09.hs
@@ -1,9 +1,10 @@
 import Data.Char
 import Data.List
+import Data.List.Split
 
 type Id = Int
 type Len = Int
-data Block = Data Id Len | Free Len deriving Show
+data Block = Data Id Len | Free Len deriving (Show, Eq)
 type HDD = [Block]
 
 getDataFree :: HDD -> Id -> [Len] -> HDD
@@ -31,16 +32,49 @@ frag1 acc (Free f:rest) (Data id d:pack)
 frag :: HDD -> HDD
 frag hdd = reverse $ frag1 [] hdd (reverse hdd)
 
-checksum :: Int -> Int -> HDD -> Int
-checksum acc _ [] = acc
-checksum acc ix (Data id len:rest) =
-  checksum (acc + (sum $ map (id*) [ix..ix+len-1])) (ix + len) rest
+replace :: Block -> HDD -> HDD
+replace d@(Data _ len) hdd =
+  let (fore:aft:[]) = splitOn [d] hdd in
+    fore ++ [Free len] ++ aft
+
+moveUp :: HDD -> HDD -> Block -> HDD
+moveUp acc (h@(Data hid _):rest) d@(Data did _)
+  | hid == did = (reverse acc) ++ (h:rest)
+  | otherwise = moveUp (h:acc) rest d
+moveUp acc (f@(Free fl):rest) d@(Data id dl)
+  | fl == dl = remove (d:acc) rest d
+  | fl > dl = remove (Free (fl - dl):d:acc) rest d
+  | otherwise = moveUp (f:acc) rest d
+  where
+    remove acc rest d = (reverse acc) ++ (replace d rest)
+
+defrag1 :: HDD -> HDD -> HDD
+defrag1 hdd []            = hdd
+defrag1 hdd (Free _:rest) = defrag1 hdd rest
+defrag1 hdd (d:rest)      = defrag1 (moveUp [] hdd d) rest
+
+defrag :: HDD -> HDD
+defrag hdd = defrag1 hdd (reverse hdd)
+
+checksum1 :: Int -> Int -> HDD -> Int
+checksum1 acc _ [] =
+  acc
+checksum1 acc ix (Free len:rest) =
+  checksum1 acc (ix + len) rest
+checksum1 acc ix (Data id len:rest) =
+  checksum1 (acc + (sum $ map (id*) [ix..ix+len-1])) (ix + len) rest
+
+checksum :: HDD -> Int
+checksum = checksum1 0 0
 
 part1 :: HDD -> Int
-part1 = (checksum 0 0) . frag
+part1 = checksum . frag
+
+part2 :: HDD -> Int
+part2 = checksum . defrag
 
 main = do
   file <- readFile "day09.input"
   let hdd = consume file
   putStrLn ("Part 1: " ++ (show $ part1 hdd))
-  --putStrLn ("Part 2: " ++ (show $ part2 hdd))
+  putStrLn ("Part 2: " ++ (show $ part2 hdd))